From 63042745210f90ab227f6e927b8f93279483f692 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Mon, 28 Sep 2026 23:11:58 +0000 Subject: [PATCH 01/20] feat(classifier): add switched classifier slots Port of granite-switch-staging PR #108 (feature/switched-classifier). Adds the switched classifier: a classifier slot is an alternative occupant of an adapter slot, holding a small linear classification head instead of a LoRA. It is fired by the same control-token machinery -- one <|name|> token, one adapter index -- but where a LoRA modifies the hidden states, a classifier reads the clean base hidden state and produces per-label logits. WHAT * Compose: classifier slots declared in a YAML manifest via `kind: classifier` and `labels:`, alongside LoRA slots in one checkpoint. `type` (token placement) and `kind` (LoRA vs classifier) are independent axes. Slots are stably reordered LoRAs-first / classifiers-last so the LoRA weight transfer stays positional and the head bank pads uniformly. * Trained-head loading with a different label count per slot: heads live in one stacked bank padded to `max_classifier_labels`, mirroring the padding rank in the LoRA bank. * Label words resolve to single token ids at compose time, failing loud and listing offenders when a word is not atomic under the base tokenizer. * Verdict exit: the request's final logit row is rewritten so the winning label's token id wins the argmax, so the endpoint emits the literal label word through the standard OpenAI-compatible server with no custom client. HF and vLLM backends kept in sync. * Chat template places the control token at the end of the last user turn, before that turn's closing marker, for both the role-marker (4.0/4.1) and ChatML (4.2) template families. * MultiSwitch: an identity `split_indices` on both backends, routing the classifier stream to zeros, since a classifier slot loads only for a quick verdict. WHY Detection/classification tasks do not need a generative adapter. A linear head over the base hidden state is far cheaper than a LoRA, and riding the existing control-token and slot machinery means it composes, serves and activates exactly like any other adapter -- no separate endpoint, no custom decode path. Tutorial: tutorials/guides/classifier_slots.md Signed-off-by: IRENEKO --- .../composer/adapter_discovery.py | 33 +- src/granite_switch/composer/arch.py | 4 +- .../composer/compose_granite_switch.py | 69 ++- src/granite_switch/composer/compose_utils.py | 123 ++++- .../composer/reporting/population_table.py | 8 +- .../composer/tokenizer_setup.py | 119 ++++- .../composer/weight_transfer.py | 148 ++++++ src/granite_switch/config.py | 78 ++++ src/granite_switch/hf/core/__init__.py | 2 + src/granite_switch/hf/core/classifier.py | 109 +++++ .../hf/modeling_granite_switch.py | 150 +++++- src/granite_switch/hf/switch/multi.py | 25 +- src/granite_switch/token_exchange.py | 44 ++ src/granite_switch/vllm/core/__init__.py | 2 + src/granite_switch/vllm/core/classifier.py | 93 ++++ .../vllm/core/lora_kernel_meta.py | 7 + .../vllm/granite_switch_model.py | 265 ++++++++++- src/granite_switch/vllm/switch/multi.py | 20 + tests/composer/test_adapter_loader.py | 42 +- tests/composer/test_classifier_compose.py | 409 ++++++++++++++++ tests/composer/test_tokenizer_setup.py | 44 ++ tests/hf/test_classifier.py | 375 +++++++++++++++ .../_classifier_chunked_prefill_worker.py | 437 ++++++++++++++++++ tests/vllm/test_classifier_chunked_prefill.py | 325 +++++++++++++ tutorials/README.md | 1 + tutorials/guides/classifier_slots.md | 297 ++++++++++++ 26 files changed, 3186 insertions(+), 43 deletions(-) create mode 100644 src/granite_switch/hf/core/classifier.py create mode 100644 src/granite_switch/vllm/core/classifier.py create mode 100644 tests/composer/test_classifier_compose.py create mode 100644 tests/hf/test_classifier.py create mode 100644 tests/vllm/_classifier_chunked_prefill_worker.py create mode 100644 tests/vllm/test_classifier_chunked_prefill.py create mode 100644 tutorials/guides/classifier_slots.md diff --git a/src/granite_switch/composer/adapter_discovery.py b/src/granite_switch/composer/adapter_discovery.py index 52ddb522..8b6e099f 100644 --- a/src/granite_switch/composer/adapter_discovery.py +++ b/src/granite_switch/composer/adapter_discovery.py @@ -108,17 +108,36 @@ def discover_adapters( def discover_adapters_from_yaml( manifest_path: str, -) -> list[tuple[str, str, str, str | None]]: +) -> tuple[list[tuple[str, str, str, str | None]], dict[str, dict]]: """Discover adapters from a YAML manifest file. Reads a YAML manifest that maps adapter names to their paths and types. + An entry can also mark itself a **classifier** slot with two extra keys: + + .. code-block:: yaml + + hallucination: + path: /path/to/head + type: lora # lora/alora technology (token placement) + kind: classifier # default "lora" when omitted + labels: [faithful, hallucinated] + + ``type`` and ``kind`` are independent axes: ``type`` is the LoRA technology + (lora vs alora, which drives token placement), ``kind`` is whether the slot + is a real LoRA or a classifier head. ``kind`` and ``labels`` are optional; + without them an entry is a plain LoRA adapter. The label words become token + ids at compose time (see :func:`tokenizer_setup.resolve_label_token_ids`). Args: manifest_path: Path to the YAML manifest file. Returns: + ``(found, classifier_meta)``. ``found`` is: List of ``(adapter_path, adapter_name, technology, source)`` tuples. The source is set to the manifest path for traceability. + ``classifier_meta`` maps name -> + ``{"kind", "labels"}`` for each entry that set ``kind``/``labels``; + names not in it are plain LoRA slots. """ import yaml @@ -126,6 +145,7 @@ def discover_adapters_from_yaml( print(f" Loading adapters manifest: {path.name}") found = [] + classifier_meta: dict[str, dict] = {} if path.is_file() and path.suffix in (".yaml", ".yml"): with open(path) as f: adapters_config = yaml.safe_load(f) @@ -137,7 +157,16 @@ def discover_adapters_from_yaml( # Use manifest path as source for traceability found.append((adapter_path, name, tech, manifest_path)) - return found + # Optional classifier slots. + kind = info.get("kind") + labels = info.get("labels") + if kind is not None or labels is not None: + classifier_meta[name] = { + "kind": kind or "lora", + "labels": labels, + } + + return found, classifier_meta def _report_module_contributions( diff --git a/src/granite_switch/composer/arch.py b/src/granite_switch/composer/arch.py index 62163433..279de384 100644 --- a/src/granite_switch/composer/arch.py +++ b/src/granite_switch/composer/arch.py @@ -112,12 +112,14 @@ class ArchDescriptor: # LoRA parameter keywords (for filtering) lora_keywords: list[str] = field(default_factory=lambda: ["lora_A", "lora_B"]) - # Non-LoRA buffer keywords to exclude from base validation + # Non-LoRA/classifier buffer keywords to exclude from base validation buffer_keywords: list[str] = field( default_factory=lambda: [ "adapter_token_ids", "adapter_scalings", "control_to_substitute_lut", + "classifier_head", + "adapter_kind_lut", ] ) diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index 8636ffa0..58aec734 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -69,6 +69,7 @@ from granite_switch.composer.compose_utils import GraniteSwitchComposer from granite_switch.composer.reporting import generate_compose_report, write_build_doc from granite_switch.composer.tokenizer_setup import ( + ANCHOR_MODE_CLASSIFIER, ANCHOR_MODE_SR, add_audio_token, add_control_tokens, @@ -77,6 +78,7 @@ configure_chat_template, find_reserved_never_emitted_token_id, load_activation_anchor, + resolve_label_token_ids, ) from granite_switch.composer.validator import validate_control_lut from granite_switch.composer.weight_transfer import validate_untied_lm_head_saved @@ -981,6 +983,7 @@ def build(): step_start = time.time() discovered_adapters = [] + classifier_meta: dict[str, dict] = {} if args.adapters: print("\n" + "=" * 80) print("Processing adapters") @@ -1022,8 +1025,9 @@ def build(): print(msg) discovered_adapters.extend(found) elif (path := Path(entry)).is_file() and path.suffix in (".yaml", ".yml"): - found = discover_adapters_from_yaml(entry) + found, manifest_classifier_meta = discover_adapters_from_yaml(entry) discovered_adapters.extend(found) + classifier_meta.update(manifest_classifier_meta) else: # Single adapter directory resolved = Path(resolved_path) @@ -1065,11 +1069,21 @@ def build(): # token lands, and SR has exactly one usable activation point — the anchor at # the end of the generation prompt — so a mislabelled SR checkpoint must not # be able to fall through to aLoRA or sequence-start placement. + # A classifier slot is labelled here for the same reason: its control token is + # an end-locator, which is a placement mode, not a weight format, and the + # manifest's ``kind`` decides it. An adapter satisfies at most one test -- an SR + # adapter carries ``cross_stream`` weights, a classifier slot carries a + # classifier head instead -- and SR goes first because it is read from the + # artifact rather than from a label. external_discovered = [ ( path, name, - ANCHOR_MODE_SR if path and is_shadow_residual_adapter(path) else tech, + ANCHOR_MODE_SR + if path and is_shadow_residual_adapter(path) + else ANCHOR_MODE_CLASSIFIER + if classifier_meta.get(name, {}).get("kind") == "classifier" + else tech, source, ) for path, name, tech, source in discovered_adapters @@ -1079,6 +1093,11 @@ def build(): ] all_discovered = external_discovered + built_in_discovered + # Order LoRA slots first, classifier slots last (Python sorts False before True) + all_discovered.sort( + key=lambda t: classifier_meta.get(t[1], {}).get("kind", "lora") == "classifier" + ) + has_external = len(external_discovered) > 0 has_built_in = len(built_in_discovered) > 0 @@ -1227,6 +1246,52 @@ def build(): _anchor_text, _sub, _mode = load_activation_anchor(_path, tokenizer) adapter_substitute_token_ids[_sub_offset + _i] = _sub + # adapter_kinds marks which slots are classifiers (same control token as a + # LoRA); resolve each slot's label words to token ids for the vLLM verdict. + if classifier_meta: + adapter_kinds = [ + classifier_meta.get(name, {}).get("kind", "lora") for name in adapter_names + ] + classifier_names = [ + name + for name in adapter_names + if classifier_meta.get(name, {}).get("kind") == "classifier" + ] + + if classifier_names: + # Full width (length num_adapters, indexed by adapter_index - 1): each + # classifier slot's label token ids. The config derives the per-slot + # counts and the padded bank width from this. + per_slot_label_names: list[list[str] | None] = [] + per_slot_label_token_ids: list[list[int] | None] = [] + for name, kind in zip(adapter_names, adapter_kinds): + if kind != "classifier": + per_slot_label_names.append(None) + per_slot_label_token_ids.append(None) + continue + labels = classifier_meta[name].get("labels") + if not labels: + raise ValueError( + f"Adapter '{name}' is declared kind=classifier but has no " + f"'labels' in the manifest. Classifier slots need label " + f"words (e.g. labels: [safe, unsafe]) so the verdict exit " + f"can emit them." + ) + label_names = list(labels) + per_slot_label_names.append(label_names) + per_slot_label_token_ids.append( + resolve_label_token_ids(tokenizer, label_names) + ) + + optional_kwargs["adapter_kinds"] = adapter_kinds + optional_kwargs["classifier_label_token_ids"] = per_slot_label_token_ids + print("\nClassifier slots:") + for name, words, ids in zip( + adapter_names, per_slot_label_names, per_slot_label_token_ids + ): + if ids is not None: + print(f" {name} → labels {words} → token ids {ids}") + model = GraniteSwitchComposer.from_base_and_adapters( base_model_name_or_path=base_model_local_path, adapter_paths=adapter_paths, diff --git a/src/granite_switch/composer/compose_utils.py b/src/granite_switch/composer/compose_utils.py index 21e3595f..b8982a1d 100644 --- a/src/granite_switch/composer/compose_utils.py +++ b/src/granite_switch/composer/compose_utils.py @@ -17,7 +17,11 @@ ) from .arch import resolve_arch from .validator import validate_all_parameters, validate_cross_stream_population -from .weight_transfer import transfer_adapter_weights, transfer_base_weights +from .weight_transfer import ( + transfer_adapter_weights, + transfer_base_weights, + transfer_classifier_weights, +) class GraniteSwitchComposer: @@ -80,6 +84,35 @@ def from_base_and_adapters( num_built_in = len(built_in_adapter_names) num_total = num_external + num_built_in + # Partition external slots by kind. + ext_kinds = kwargs.get("adapter_kinds") + if ext_kinds is None: + ext_kinds = ["lora"] * num_external + else: + ext_kinds = list(ext_kinds)[:num_external] + lora_slots = [i for i in range(num_external) if ext_kinds[i] != "classifier"] + classifier_slots = [ + i for i in range(num_external) if ext_kinds[i] == "classifier" + ] + + # Compose orders LoRAs first, classifiers last, so classifier slots must + # form a contiguous suffix (LoRAs occupy 0..k-1, classifiers k..). + if classifier_slots and classifier_slots != list( + range(lora_slots[-1] + 1 if lora_slots else 0, num_external) + ): + raise ValueError( + "Classifier slots must form a contiguous suffix (all LoRA slots " + f"first, then classifiers); got kinds {ext_kinds}. Order the " + "adapters LoRAs-first (compose does this automatically)." + ) + lora_adapter_paths = [adapter_paths[i] for i in lora_slots] + classifier_adapter_paths = [adapter_paths[i] for i in classifier_slots] + lora_adapter_names = ( + [adapter_names[i] for i in lora_slots] + if adapter_names is not None + else None + ) + # --- Step 1: Resolve architecture --- # Pre-scan for cross_stream (SR dual-stream) to select the right arch. print(f"Loading config from {base_model_name_or_path}...") @@ -140,16 +173,28 @@ def from_base_and_adapters( f"(cross_stream_rank={cross_stream_rank})" ) - if adapter_paths: - lora_rank, _lora_alpha, adapter_ranks, adapter_alphas = detect_lora_config( - adapter_paths + # adapter_ranks/adapter_alphas stay full external width (one entry per + # slot, classifiers included) so config.adapter_ranks is length + # num_adapters. A classifier slot's LoRA row never fires: split_indices + # zeroes the LoRA stream at classifier positions. + if lora_adapter_paths: + lora_rank, lora_alpha, lora_ranks, lora_alphas = detect_lora_config( + lora_adapter_paths ) lora_target_modules, source_analysis = detect_present_modules( - adapter_paths, + lora_adapter_paths, arch, - adapter_names=adapter_names, + adapter_names=lora_adapter_names, ) + # LoRAs lead, classifiers follow, so the full lists are a concat: + # detected LoRA values, then a placeholder per classifier slot. The + # placeholder is the detected max rank, keeping max(adapter_ranks) == + # max_lora_rank. + num_ext_classifier = num_external - len(lora_slots) + adapter_ranks = list(lora_ranks) + [lora_rank] * num_ext_classifier + adapter_alphas = list(lora_alphas) + [float(lora_rank)] * num_ext_classifier + # Extend adapter_ranks with built-in entries if num_built_in > 0: if built_in_lora_rank != lora_rank: @@ -162,21 +207,22 @@ def from_base_and_adapters( list(adapter_ranks) + [built_in_lora_rank] * num_built_in ) lora_rank = max(lora_rank, built_in_lora_rank) + elif classifier_adapter_paths or num_built_in > 0: + # No LoRA slots but the model still has slots (classifier and/or + # built-in). Give every slot a placeholder LoRA rank so the + # (unused) LoRA bank is well-formed and config validation passes. + lora_rank = built_in_lora_rank + adapter_ranks = [built_in_lora_rank] * num_total + adapter_alphas = {} + lora_target_modules = None + source_analysis = {} else: - # Built-in only or zero-adapter skinning - if num_built_in > 0: - lora_rank = built_in_lora_rank - adapter_ranks = [built_in_lora_rank] * num_built_in - adapter_alphas = {} - # Auto-detect lora_target_modules (attention-only switch model) - lora_target_modules = None - source_analysis = {} - else: - lora_rank = 0 - adapter_ranks = None - adapter_alphas = {} - lora_target_modules = [] - source_analysis = {} + # Zero-adapter skinning (base model only). + lora_rank = 0 + adapter_ranks = None + adapter_alphas = {} + lora_target_modules = [] + source_analysis = {} # --- Step 4: Build switch config from arch descriptor --- # Copy config fields driven by architecture descriptor @@ -261,10 +307,20 @@ def from_base_and_adapters( base_model_name_or_path, model, switch_config, arch ) - if adapter_paths: - # --- Step 7: Transfer adapter weights --- + adapter_mapping = {} + classifier_mapping = {} + + if lora_adapter_paths: + # --- Step 7: Transfer LoRA adapter weights --- + # LoRAs are the contiguous leading slots, so the leading k entries of + # adapter_ranks/adapter_alphas are exactly the LoRA slots' values and + # each path's list position is its bank row (no slot remapping). + n_lora = len(lora_adapter_paths) adapter_mapping = transfer_adapter_weights( - adapter_paths, model, adapter_alphas, arch + lora_adapter_paths, + model, + list(adapter_alphas)[:n_lora], + arch, ) # --- Step 8: Validate --- @@ -273,8 +329,12 @@ def from_base_and_adapters( validate_all_parameters( model, arch, - adapter_paths=adapter_paths, - adapter_names=adapter_names[:num_external], + adapter_paths=lora_adapter_paths, + adapter_names=( + [adapter_names[i] for i in lora_slots] + if adapter_names is not None + else None + ), target_module_sets=target_module_sets, ) if dual_stream: @@ -282,6 +342,18 @@ def from_base_and_adapters( else: adapter_mapping = {} + if classifier_adapter_paths: + # --- Step 7b: Transfer trained classifier-head weights --- + # Each slot's head is validated against its own label count and + # copied into the top rows of the padded bank. + classifier_mapping = transfer_classifier_weights( + classifier_adapter_paths, + classifier_slots, + model, + num_labels_per_slot=switch_config.classifier_num_labels_per_slot, + hidden_size=switch_config.hidden_size, + ) + print("\nModel created successfully!") print(f" Base model: {base_model_name_or_path}") print( @@ -295,6 +367,7 @@ def from_base_and_adapters( model._build_mappings = { "base": base_mapping, "adapter": adapter_mapping, + "classifier": classifier_mapping, "source_analysis": source_analysis, # Per-external-adapter alpha, parallel to adapter_paths. When no # external adapters are provided, detect_lora_config isn't called diff --git a/src/granite_switch/composer/reporting/population_table.py b/src/granite_switch/composer/reporting/population_table.py index 0e94da46..5cfec913 100644 --- a/src/granite_switch/composer/reporting/population_table.py +++ b/src/granite_switch/composer/reporting/population_table.py @@ -49,12 +49,16 @@ def generate_adapter_population_table( if target_module_sets is None: target_module_sets = load_adapter_target_modules(adapter_paths) + # target_module_sets covers LoRA slots only (classifier slots carry no LoRA + # weights, so source analysis never produced a target set for them). adapter_configs = [ { "rank": adapter_ranks[i] if adapter_ranks else max_rank, - "target_modules": target_module_sets[i], + "target_modules": target_module_sets[i] + if i < len(target_module_sets) + else set(), } - for i in range(len(adapter_paths)) + for i in range(num_adapters) ] # Define module types (rows) — only modules in lora_target_modules diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 29be363b..6e9cac57 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -145,6 +145,10 @@ def detect_template_format(chat_template: str | None) -> TemplateFormat | None: #: the end of the generation prompt (Shadow Residual, ``last_context_token``). ANCHOR_MODE_SR = "sr" +#: Placement mode where the control token is emitted *after* the last content +#: token of the last user turn (classifier slot, manifest ``kind``). +ANCHOR_MODE_CLASSIFIER = "classifier" + def _load_alora_invocation_token_ids(adapter_path: str) -> list[int]: """Load alora_invocation_tokens from adapter_config.json. @@ -864,6 +868,30 @@ def build_substitute_token_ids( return substitute_ids +def resolve_label_token_ids(tokenizer, label_names: list[str]) -> list[int]: + """Turn each classifier label word into its single token id.""" + token_ids: list[int] = [] + non_atomic: list[str] = [] + for word in label_names: + # Same encode as the control-token check: bare word, no special tokens. + encoded = tokenizer.encode(word, add_special_tokens=False) + if len(encoded) != 1: + sub_tokens = tokenizer.convert_ids_to_tokens(encoded) + non_atomic.append(f"{word!r} → {len(encoded)} tokens: {sub_tokens}") + else: + token_ids.append(encoded[0]) + + if non_atomic: + raise ValueError( + f"{len(non_atomic)} classifier label word(s) do not encode to a " + f"single token, but the vLLM classifier exit writes the verdict at " + f"one token id per label. Choose a single-token synonym for each:\n " + + "\n ".join(non_atomic) + ) + + return token_ids + + def configure_chat_template( tokenizer, discovered_adapters: list[tuple[str | None, str, str, str | None]], @@ -935,6 +963,11 @@ def configure_chat_template( # still costing compute. Callers therefore classify SR from the checkpoint # weights (``is_shadow_residual_adapter``) and pass ``sr`` here, rather than # trusting a directory label that no SR training run produces. + # A classifier slot arrives as technology ``classifier``, like SR arrives as + # ``sr``: the caller derives the placement mode (SR from the checkpoint + # weights, classifier from the manifest's ``kind``) and passes it in + # the discovered tuple. ``technology`` is the placement axis this function + # dispatches on; a classifier is a fourth mode alongside lora / alora / sr. adapter_mapping: dict[str, dict[str, str]] = {} sr_anchors: set[str] = set() for adapter_info in discovered_adapters: @@ -983,6 +1016,9 @@ def configure_chat_template( mapping_entries = [] for adapter_name, info in adapter_mapping.items(): + # ``type`` is the placement mode: lora (sequence start) / alora (before + # the invocation text) / sr (replaces the generation-prompt anchor) / + # classifier (end-locator after the last content token). if "invocation_text" in info: mapping_entries.append( f" '{adapter_name}': {{'token': '{info['token']}', " @@ -1141,6 +1177,31 @@ def configure_chat_template( {%- endif %} """ + # Classifier placement: the marker goes after the last user turn's content and + # before that turn's closing marker (``<|end_of_text|>`` / ``<|im_end|>``), so + # the read point ``marker - 1`` is the position the head was trained on. + classifier_scan = ( + """{#- Classifier scan: find the last user message (the turn being classified). -#} +{%- if ns.adapter_token and ns.adapter_type == 'classifier' %} + {%- for _msg in """ + + fmt.loop_var_source + + """ %} + {%- if _msg.role == 'user' %} + {%- set ns.classifier_target_idx = loop.index0 %} + {%- endif %} + {%- endfor %} +{%- endif %} +""" + ) + # In-loop emit: at the target user message, emit the marker right before the + # turn-closing marker. ``{{- ns.adapter_token }}`` prints the control token; + # the surrounding content and turn-close are emitted by the base template. + classifier_emit = ( + "{%- if ns.adapter_token and ns.adapter_type == 'classifier' " + "and loop.index0 == ns.classifier_target_idx %}" + "{{- ns.adapter_token }}{%- endif %}" + ) + # Build the modified template modified_chat_template = adapter_map_def + adapter_lookup + base_chat_template @@ -1165,6 +1226,7 @@ def configure_chat_template( "\n adapter_invocation_text=adapter_invocation_text," "\n adapter_invocation_drop_chars=adapter_invocation_drop_chars," "\n alora_target_idx=-1," + "\n classifier_target_idx=-1," "\n skip_next_role_marker=false" "\n )" ) @@ -1251,6 +1313,7 @@ def configure_chat_template( modified_chat_template = ( modified_chat_template[:insertion_point] + alora_pass1 + + classifier_scan + modified_chat_template[insertion_point:] ) @@ -1282,6 +1345,58 @@ def configure_chat_template( + modified_chat_template[insertion_point:] ) + # Inject the classifier emit at the end of the last user turn, before that + # turn's closing marker. Format-specific because the two families emit the + # turn boundary differently (granite_format fuses content and the close in + # one emission; chatml emits the close as a standalone statement). + classifier_names = sorted( + n for n, i in adapter_mapping.items() if i["type"] == ANCHOR_MODE_CLASSIFIER + ) + if classifier_names and fmt.name == "chatml": + # Emit right before the user/system branch's standalone + # ``{{- '<|im_end|>\n' }}``. That close follows the low-effort if/else, + # so anchor on ``{%- endif %}\n{{- '<|im_end|>\n' }}`` scoped to the + # user branch: the FIRST ``<|im_end|>`` emission that is preceded by an + # ``{%- endif %}`` is the user/system one (assistant/tool closes are not + # preceded by an if/else over content). + chatml_user_close_re = re.compile( + r"(\{%-\s*endif\s*%\}\s*)(\{\{-\s*'<\|im_end\|>\\n'\s*\}\})" + ) + m = chatml_user_close_re.search(modified_chat_template) + if m: + insertion_point = m.start(2) + modified_chat_template = ( + modified_chat_template[:insertion_point] + + classifier_emit + + modified_chat_template[insertion_point:] + ) + elif classifier_names: + # granite_format: split the fused user emission so the marker prints + # after ``content.val`` and before ``'<|end_of_text|>\n'``. Rewrite + # ``... + content.val + '<|end_of_text|>\n' }}`` into + # ``... + content.val }}{{- '<|end_of_text|>\n' }}``. + gf_user_emit = ( + "'<|start_of_role|>' + message.role + '<|end_of_role|>' " + "+ content.val + '<|end_of_text|>\\n'" + ) + gf_user_emit_split = ( + "'<|start_of_role|>' + message.role + '<|end_of_role|>' " + "+ content.val }}" + classifier_emit + "{{- '<|end_of_text|>\\n'" + ) + if gf_user_emit in modified_chat_template: + modified_chat_template = modified_chat_template.replace( + gf_user_emit, gf_user_emit_split, 1 + ) + + if classifier_names and classifier_emit not in modified_chat_template: + raise ValueError( + f"Classifier slot(s) {classifier_names} were requested, but the " + f"classifier control-token emit could not be injected into the " + f"{fmt.name} chat template: the user-turn closing anchor was not found. " + f"Update the anchor in configure_chat_template to match this base's " + f"template." + ) + # Insert alora fallback before generation prompt gen_prompt_pattern = r"(\{%- if add_generation_prompt %\})" match = re.search(gen_prompt_pattern, modified_chat_template) @@ -1385,7 +1500,9 @@ def _split_merged(match: "re.Match") -> str: tokenizer.chat_template = modified_chat_template print(f"Chat template configured with {len(adapter_mapping)} adapter mappings:") for adapter_name, info in adapter_mapping.items(): - if info["type"] == ANCHOR_MODE_SR: + if info["type"] == ANCHOR_MODE_CLASSIFIER: + placement = "at the end of the last user turn (end-locator; read point = marker - 1)" + elif info["type"] == ANCHOR_MODE_SR: placement = f"replacing '{sr_site}' in the generation prompt" anchor = next(iter(sr_anchors)) if sr_site != anchor: diff --git a/src/granite_switch/composer/weight_transfer.py b/src/granite_switch/composer/weight_transfer.py index cff280db..f5f2472c 100644 --- a/src/granite_switch/composer/weight_transfer.py +++ b/src/granite_switch/composer/weight_transfer.py @@ -15,6 +15,11 @@ from .arch import ArchDescriptor +# Filename + keys of a trained classifier-head artifact inside a classifier +# adapter directory. +CLASSIFIER_HEAD_FILE = "classifier_head.safetensors" +CLASSIFIER_HEAD_KEYS = ("weight", "bias") + # --------------------------------------------------------------------------- # Base weight loading # --------------------------------------------------------------------------- @@ -824,3 +829,146 @@ def validate_untied_lm_head_saved(output_path, expected_vocab_size, hidden_size) f" Untied LM head validated: lm_head.weight present, shape " f"[{expected_vocab_size}, {hidden_size}]" ) + + +# --------------------------------------------------------------------------- +# Classifier head transfer +# --------------------------------------------------------------------------- + + +def transfer_classifier_weights( + classifier_paths: list[str], + slot_indices: list[int], + model, + num_labels_per_slot: list[int], + hidden_size: int, + return_mapping: bool = True, +) -> dict | None: + """Load trained classifier heads and transfer them into the switch model. + + A classifier slot is the classifier-head analog of a LoRA adapter: it shares + the same control-token / adapter-index space, so its head lives in bank row + ``slot = adapter_index - 1`` of ``model.model.classifier_head``. Each + classifier directory holds a ``classifier_head.safetensors`` with a + ``weight`` ``[n_i, hidden]`` and a ``bias`` ``[n_i]`` where ``n_i`` is this + slot's own label count. The head bank is padded to ``max_classifier_labels``. + + Args: + classifier_paths: Directories of the classifier slots (classifier-only). + slot_indices: Global bank row for each path (``adapter_index - 1``), + aligned with *classifier_paths*. Classifiers are the non-zero-based + ``k+1..n`` suffix, so these are not ``range(len)`` and must be passed. + model: GraniteSwitch model whose ``classifier_head`` is filled in place. + num_labels_per_slot: Full-width per-slot label counts + (``config.classifier_num_labels_per_slot``, indexed by + ``adapter_index - 1``). + hidden_size: Expected hidden size (``config.hidden_size``). + return_mapping: If True, return detailed mapping information. + + Returns: + Mapping record dict, or None if *return_mapping* is False. + """ + from safetensors.torch import load_file + + head = getattr(model.model, "classifier_head", None) + if head is None: + raise ValueError( + "transfer_classifier_weights was given classifier paths " + f"{classifier_paths} but the model built no classifier_head. This " + "means adapter_kinds never reached the config (no slot is " + "kind=classifier). Pass adapter_kinds with a 'classifier' entry so " + "the head is constructed before loading trained weights." + ) + + _validate_classifier_transfer(model, slot_indices) + + mapping_record = {"source_params": [], "target_params": [], "mappings": []} + + print(f"Loading {len(classifier_paths)} trained classifier head(s)...") + with torch.no_grad(): + for path, slot in zip(classifier_paths, slot_indices): + n_labels = num_labels_per_slot[slot] + head_file = Path(path) / CLASSIFIER_HEAD_FILE + if not head_file.exists(): + raise FileNotFoundError( + f"Classifier slot '{path}' is missing its trained head " + f"'{CLASSIFIER_HEAD_FILE}'. A classifier directory must " + f"contain {CLASSIFIER_HEAD_FILE} with keys " + f"{CLASSIFIER_HEAD_KEYS} (weight [n_labels, hidden], " + f"bias [n_labels]); it has no adapter_model.safetensors." + ) + + sd = load_file(str(head_file)) + keys = set(sd.keys()) + if keys != set(CLASSIFIER_HEAD_KEYS): + raise ValueError( + f"{head_file} must contain exactly keys " + f"{set(CLASSIFIER_HEAD_KEYS)}, got {keys}." + ) + + weight, bias = sd["weight"], sd["bias"] + if tuple(weight.shape) != (n_labels, hidden_size): + raise ValueError( + f"{head_file}: weight shape {tuple(weight.shape)} does not " + f"match this slot's expected (n_labels, hidden) = " + f"({n_labels}, {hidden_size}). Check the head's label count " + f"against classifier_num_labels_per_slot[{slot}]={n_labels}." + ) + if tuple(bias.shape) != (n_labels,): + raise ValueError( + f"{head_file}: bias shape {tuple(bias.shape)} does not " + f"match this slot's expected (n_labels,) = ({n_labels},)." + ) + + # Copy into the top n_labels rows of the padded bank row; the rest + # stay zero. + dst_w = head.weight[slot] + head.weight[slot, :n_labels, :].copy_( + weight.to(dtype=dst_w.dtype, device=dst_w.device) + ) + head.bias[slot, :n_labels].copy_( + bias.to(dtype=dst_w.dtype, device=dst_w.device) + ) + + print(f" Classifier: {path} -> classifier_head[{slot}, :{n_labels}]") + mapping_record["source_params"].append(f"{path}::weight") + mapping_record["source_params"].append(f"{path}::bias") + mapping_record["mappings"].append( + { + "source": [f"{path}::weight", f"{path}::bias"], + "target": f"model.classifier_head[{slot}]", + "type": "classifier", + } + ) + + mapping_record["target_params"] = [ + f"model.classifier_head.weight[{s}]" for s in slot_indices + ] + [f"model.classifier_head.bias[{s}]" for s in slot_indices] + + print(f"Loaded {len(classifier_paths)} classifier head(s)") + return mapping_record if return_mapping else None + + +def _validate_classifier_transfer(model, slot_indices: list[int]): + """Sanity-check the classifier slots before transfer. + + Guards against two silent-corruption modes: two directories claiming the + same bank row, and a slot index outside the head's ``[num_adapters, ...]`` + bank. + """ + if len(slot_indices) != len(set(slot_indices)): + raise ValueError( + f"Duplicate classifier slot indices {slot_indices}; each classifier " + "must map to a distinct bank row." + ) + num_rows = model.model.classifier_head.weight.shape[0] + for slot in slot_indices: + if not (0 <= slot < num_rows): + raise ValueError( + f"Classifier slot {slot} is out of range for a classifier bank " + f"of {num_rows} rows (num_adapters)." + ) + print( + f"Classifier transfer: {len(slot_indices)} slot(s) " + f"{sorted(slot_indices)} into a {num_rows}-row bank" + ) diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index 9ee20816..37e795c2 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -181,6 +181,9 @@ def __init__( # Shadow Residual (SR) parameters cross_stream_rank: int | None = None, dual_stream: bool = False, + # Classifier substitution parameters + adapter_kinds: list[str] | None = None, + classifier_label_token_ids: list[list[int] | None] | None = None, # vLLM residual-norm convention (for bit-exact skinning equivalence) fused_add_norm: bool = False, # Parent class defaults (Granite 4 dense configuration) @@ -381,6 +384,81 @@ def __init__( self.max_lora_rank = max_lora_rank self.adapter_ranks = adapter_ranks + # Each adapter slot is either a LoRA adapter ("lora", the default) or + # a classifier head ("classifier"). A classifier slot still uses a + # control token and an adapter index; the switch routes its token to + # that index like lora adapter, but the index is tagged as a classifier + # so the LoRA path no-ops there and the classifier head runs instead. + _VALID_KINDS = ("lora", "classifier") + if adapter_kinds is None: + adapter_kinds = ["lora"] * num_adapters + else: + if len(adapter_kinds) != num_adapters: + raise ValueError( + f"adapter_kinds length ({len(adapter_kinds)}) must equal " + f"num_adapters ({num_adapters})." + ) + invalid = [k for k in adapter_kinds if k not in _VALID_KINDS] + if invalid: + raise ValueError( + f"adapter_kinds values must be one of {_VALID_KINDS}; " + f"got invalid entries {invalid}" + ) + self.adapter_kinds = adapter_kinds + + # Classifier metadata resolved at compose time. + if classifier_label_token_ids is not None: + if len(classifier_label_token_ids) != num_adapters: + raise ValueError( + f"classifier_label_token_ids length " + f"({len(classifier_label_token_ids)}) must equal num_adapters " + f"({num_adapters}); it is indexed by adapter_index - 1, with " + f"None on non-classifier slots." + ) + for i, (kind, ids) in enumerate( + zip(adapter_kinds, classifier_label_token_ids) + ): + if kind == "classifier": + if not ids: + raise ValueError( + f"classifier slot {i} must declare its label token ids; " + f"got {ids!r}." + ) + if any(tid < 0 for tid in ids): + raise ValueError( + f"classifier_label_token_ids[{i}] must all be >= 0 " + f"(real token ids); got {ids}" + ) + elif ids: + raise ValueError( + f"non-classifier slot {i} must have no label token ids; " + f"got {ids!r}." + ) + elif any(k == "classifier" for k in adapter_kinds): + raise ValueError( + "adapter_kinds declares a classifier slot but " + "classifier_label_token_ids is unset. The classifier exit emits its " + "verdict as a label word, so the per-label token ids are required. " + "Re-compose so config.json carries classifier_label_token_ids." + ) + self.classifier_label_token_ids = classifier_label_token_ids + + # Derived: per-slot label counts (0 on LoRA slots) and the padded bank + # width the classifier head is built with. + self.classifier_num_labels_per_slot = [ + len(ids) if ids else 0 for ids in (classifier_label_token_ids or []) + ] or [0] * num_adapters + self.max_classifier_labels = max(self.classifier_num_labels_per_slot, default=0) + + # Derived: the control-token ids that mark a classifier slot. Both backends + # locate a verdict's read point by matching these in the original input_ids, + # so the id set is defined once here. + self.classifier_control_token_ids = [ + tid + for tid, kind in zip(adapter_token_ids or [], adapter_kinds) + if kind == "classifier" + ] + # Default LoRA target module groups. # Dynamically determined based on model architecture. # Empty when num_adapters == 0 (no LoRA to apply). diff --git a/src/granite_switch/hf/core/__init__.py b/src/granite_switch/hf/core/__init__.py index 38341501..28e7e412 100644 --- a/src/granite_switch/hf/core/__init__.py +++ b/src/granite_switch/hf/core/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 """Core LoRA primitives for Granite Switch (HuggingFace).""" +from .classifier import SwitchedClassifierHead from .lora import ( GraniteLoRAEmbeddedAttention, MergedSwitchedLoRALinear, @@ -10,5 +11,6 @@ __all__ = [ "GraniteLoRAEmbeddedAttention", "MergedSwitchedLoRALinear", + "SwitchedClassifierHead", "SwitchedLoRALinear", ] diff --git a/src/granite_switch/hf/core/classifier.py b/src/granite_switch/hf/core/classifier.py new file mode 100644 index 00000000..929bf1d0 --- /dev/null +++ b/src/granite_switch/hf/core/classifier.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Switched classifier heads for Granite Switch. + +A classifier head is an optional substitute for a LoRA adapter: an adapter slot +is either a LoRA adapter or a classifier head, selected by the same control +token and adapter index. Where a classifier is selected, the LoRA path no-ops +and the classifier reads clean base-model hidden states. + +Classifier positions are expected to be sparse, and the weight bank is replicated +on every rank at this point; revisit if classifiers ever cover most tokens or +num_labels grows large. +""" + +import torch +import torch.nn as nn + + +class SwitchedClassifierHead(nn.Module): + """Per-token classifier heads selected by classifier indices. + + Stacked weight bank indexed by ``(index - 1)``, like ``SwitchedLoRALinear``. + Index ``0`` = no classifier (base); ``1+`` selects that slot's head. + Positions not selected are left at zero. + + Args: + hidden_size: Input feature dimension (the model hidden size). + num_classifier_slots: ``num_adapters`` (the full index space), so + ``index - 1`` addresses the bank directly like the LoRA bank. Slots + that are not classifiers simply never get selected. + max_num_labels: Padded label width of the bank — ``max`` over the + per-slot label counts. Slots with fewer labels leave their extra + rows zero. The verdict exit slices each slot back to its real label + count. + """ + + def __init__( + self, + hidden_size: int, + num_classifier_slots: int, + max_num_labels: int, + ): + super().__init__() + self.hidden_size = hidden_size + self.num_classifier_slots = num_classifier_slots + self.num_labels = max_num_labels + + # Slot s (0-based) serves adapter index s + 1, matching the LoRA bank's + # tensor_idx = adapter_idx - 1. + # + # A frozen inference bank whose real values are copied in from a trained + # head at compose time. Zero value for both an unloaded slot's row and a + # loaded slot's pad rows (label counts below max_num_labels). + self.weight = nn.Parameter( + torch.zeros(num_classifier_slots, max_num_labels, hidden_size) + ) + self.bias = nn.Parameter(torch.zeros(num_classifier_slots, max_num_labels)) + + def forward( + self, + x: torch.Tensor, + classifier_indices: torch.Tensor, + ) -> torch.Tensor: + """Compute per-token classifier logits for selected positions. + + Args: + x: Input tensor [batch, seq_len, hidden_size] or [num_tokens, hidden_size]. + classifier_indices: Per-token classifier selection [batch, seq_len] + or [num_tokens]; 0 = no classifier, 1+ = slot index. + + Returns: + Logits with the selected slot's head applied per token: + [batch, seq_len, num_labels] or [num_tokens, num_labels]. Positions + with index 0 (and any not selected by this bank) are left at zero. + """ + # Flatten for token-level processing (same as SwitchedLoRALinear). + original_shape = x.shape + if x.dim() == 3: + batch_size, seq_len, _ = x.shape + x_flat = x.reshape(-1, self.hidden_size) + else: + x_flat = x + + logits_flat = x_flat.new_zeros((x_flat.shape[0], self.num_labels)) + + idx_flat = classifier_indices.reshape(-1) + + # Per-token gather/scatter loop (same as SwitchedLoRALinear.forward) + mask = idx_flat > 0 + if mask.any(): + active = idx_flat[mask].unique() + for adapter_idx in active: + token_mask = idx_flat == adapter_idx + token_indices = torch.where(token_mask)[0] + + if len(token_indices) == 0: + continue + + slot = adapter_idx - 1 # 1-indexed → 0-indexed, like the LoRA bank + w = self.weight[slot] # [num_labels, hidden_size] + b = self.bias[slot] # [num_labels] + + x_sel = x_flat[token_indices] # [num_sel, hidden_size] + # Write (not add): a classifier replaces the LoRA. + logits_flat[token_indices] = x_sel @ w.t() + b + + # Reshape back if needed. + if len(original_shape) == 3: + return logits_flat.view(batch_size, seq_len, self.num_labels) + return logits_flat diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index 22d5c4c7..207c75b6 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -29,6 +29,7 @@ from granite_switch.config import GraniteSwitchConfig +from .core.classifier import SwitchedClassifierHead from .core.lora import ( GraniteLoRAEmbeddedAttention, SwitchedLoRALinear, @@ -38,6 +39,7 @@ logger = logging.get_logger(__name__) + # transformers 5.9.0 renamed `input_embeds` -> `inputs_embeds` and dropped the # unused `cache_position` kwarg in `create_causal_mask`. _TRANSFORMERS_GE_5_9 = _parse_version(transformers.__version__) >= _parse_version( @@ -417,6 +419,21 @@ def __init__(self, config: GraniteSwitchConfig): self.switch = None self.adapter_token_ids = None + # Classifier heads: optional alternative for LoRA adapters. A classifier + # slot shares the adapter index/token machinery, but its index is tagged + # as "classifier" (via the switch's adapter_kind_lut). The switch's + # split_indices() routes those positions out of the LoRA stream (so LoRA + # no-ops there) and into a classifier stream, which this head consumes + # from the final post-norm hidden state. + if "classifier" in getattr(config, "adapter_kinds", []): + self.classifier_head = SwitchedClassifierHead( + hidden_size=config.hidden_size, + num_classifier_slots=config.num_adapters, + max_num_labels=config.max_classifier_labels, + ) + else: + self.classifier_head = None + # Decoder layers if config.num_adapters > 0: layer_offset = self.switch.num_cache_layers @@ -587,6 +604,17 @@ def forward( device=device, ) + # Split the single adapter_indices stream into a LoRA stream and + # a classifier stream. Classifier positions are zeroed in lora_indices + # (so the LoRA path no-ops there). + if self.switch is not None: + lora_indices, classifier_indices = self.switch.split_indices( + adapter_indices + ) + else: + lora_indices = adapter_indices + classifier_indices = torch.zeros_like(adapter_indices) + # Embed once, on the (possibly-rewritten) input_ids. The decoder is # token-exchange-agnostic — it just embeds whatever the switch # passed through. @@ -594,8 +622,11 @@ def forward( inputs_embeds = self.embed_tokens(modified_input_ids) inputs_embeds = inputs_embeds * self.embedding_multiplier - # Expose adapter_indices for tests and debugging. + # Expose adapter_indices for tests and debugging. _last_adapter_indices is + # the full stream (LoRA + classifier); _last_classifier_indices is the + # classifier-only stream the verdict reads. self._last_adapter_indices = adapter_indices + self._last_classifier_indices = classifier_indices position_embeddings = None if self.rotary_emb is not None: @@ -628,7 +659,7 @@ def forward( use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, - adapter_indices=adapter_indices, + adapter_indices=lora_indices, **kwargs, ) @@ -646,6 +677,80 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) + # Classifier heads read the final post-norm hidden state (the same + # representation the LM head consumes) and emit one verdict per request + # when routed to a classifier. A detect request is prefill-only: it + # classifies the prompt and ships the verdict as the replacements to + # rewrite that request's logit row into its label word. Reading the + # final state is the initial supported read point; earlier layers can + # be added later. + + # Per-request gate so a mixed model still serves LoRA/base generation. + active_classifier = self.classifier_head is not None and bool( + (classifier_indices > 0).any() + ) + if active_classifier: + if cache_position[0] != 0: + raise RuntimeError( + "Classifier requests are detect-only (prefill, no " + "generation); got a cached decode step " + f"(cache_position starts at {int(cache_position[0])})." + ) + # The marker is the LAST classifier control token in the original + # input_ids: ``classifier_indices`` is nonzero from any classifier + # marker onward, so with several classifier slots the index alone does + # not say which slot's marker a position belongs to. A multi-turn prompt + # can also carry markers from earlier turns. + if input_ids is None: + raise RuntimeError( + "A classifier request needs input_ids to locate its control " + "token, but only inputs_embeds was provided. Classifier slots " + "classify a text prompt; pass input_ids." + ) + marker_ids = torch.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + is_marker = torch.isin(input_ids, marker_ids) # [batch, seq_len] + has_marker = is_marker.any(dim=1) + # One verdict per request: the exit rewrites a single logit row, so a + # request naming two different classifier slots has no way to report + # both. Repeated markers of the SAME slot are fine -- a multi-turn + # prompt carries them from earlier turns and the last one wins. + marked = torch.where(is_marker, input_ids, -1) # [batch, seq_len] + for row in range(batch_size): + distinct = torch.unique(marked[row]) + distinct = distinct[distinct >= 0] + if distinct.numel() > 1: + raise RuntimeError( + "A classifier request may name only one classifier slot; " + f"request {row} carries control tokens " + f"{sorted(int(t) for t in distinct)}. The verdict exit " + "rewrites one logit row, so only one slot can report." + ) + pos = torch.arange(seq_length, device=input_ids.device) + marker_idx = torch.where(is_marker, pos, -1).max(dim=1).values.clamp(min=0) + if bool((has_marker & (marker_idx == 0)).any()): + raise RuntimeError( + "Classifier marker is at position 0, so there is no last " + "content token to read (the verdict is read at marker - 1). " + "The chat template emits the marker after the last token." + ) + # A row with a marker reads its predecessor. A row without one (a plain + # LM request in a mixed batch) has no read point at all; it is parked on + # index 0 to keep the gather in range and carries slot 0, so the exit's + # is_classifier gate discards its zero verdict. + read_idx = torch.where(has_marker, marker_idx - 1, 0) + batch_ar = torch.arange(batch_size, device=hidden_states.device) + req_hidden = hidden_states[batch_ar, read_idx, :] # [batch, hidden] + req_slot = classifier_indices[batch_ar, marker_idx] # [batch] + self._last_classifier_logits = self.classifier_head(req_hidden, req_slot) + self._last_req_classifier_indices = req_slot + else: + self._last_classifier_logits = None + self._last_req_classifier_indices = None + if not return_dict: return tuple( v @@ -764,6 +869,10 @@ def forward( logits = self.lm_head(hidden_states[:, slice_indices, :]) logits = logits / self.config.logits_scaling + # Classifier exit: a classifier request reports its verdict as a + # generated label word (like a guardian LoRA). + logits = self._apply_classifier_verdict(logits) + loss = None if labels is not None: loss = self.loss_function( @@ -780,3 +889,40 @@ def forward( hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) + + def _apply_classifier_verdict(self, logits: torch.Tensor) -> torch.Tensor: + """Rewrite classifier requests' last logit row to emit their label word. + + For each classifier request set the whole vocab to ``-inf`` at the last + prompt position, place the per-label verdict scores on each label word's + token id (``classifier_label_token_ids``) so decoding emits that word. + Non-classifier rows keep their normal LM logits. + """ + verdict = getattr(self.model, "_last_classifier_logits", None) + req_slot = getattr(self.model, "_last_req_classifier_indices", None) + if verdict is None or req_slot is None: + return logits + + per_slot_label_ids = self.config.classifier_label_token_ids + + if verdict.shape[0] != logits.shape[0]: + raise RuntimeError( + f"classifier verdict rows ({verdict.shape[0]}) do not match " + f"logits rows ({logits.shape[0]}); per-request alignment broke." + ) + + is_classifier = req_slot > 0 # [batch] + if not bool(is_classifier.any()): + return logits + + # Each request writes only its own slot's label ids and only that slot's + # real label count (verdict[:, :n]); padded columns are never read. + rows = is_classifier.nonzero(as_tuple=True)[0] # [num_classifier_reqs] + logits[rows, -1, :] = float("-inf") + for r in rows.tolist(): + slot = int(req_slot[r]) + ids = per_slot_label_ids[slot - 1] + n = len(ids) + label_ids = torch.tensor(ids, dtype=torch.long, device=logits.device) + logits[r, -1, label_ids] = verdict[r, :n].to(logits.dtype) + return logits diff --git a/src/granite_switch/hf/switch/multi.py b/src/granite_switch/hf/switch/multi.py index 2212ffce..f98e5773 100644 --- a/src/granite_switch/hf/switch/multi.py +++ b/src/granite_switch/hf/switch/multi.py @@ -66,7 +66,12 @@ from transformers.cache_utils import Cache from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from ...token_exchange import apply_token_exchange, build_control_to_substitute_lut +from ...token_exchange import ( + apply_token_exchange, + build_adapter_kind_lut, + build_control_to_substitute_lut, + split_adapter_indices, +) from .codes import KerdockDGCodeGenerator, recover_count_from_signal # Large negative finite value for masking. NOT literal -inf because IEEE 754 @@ -275,6 +280,15 @@ def __init__( # non-persistent buffer is zeroed by ``from_pretrained``. This LUT uses # -1 as the "not a control token" sentinel, so an all-zero LUT would # rewrite EVERY token id to 0 in apply_token_exchange. + # adapter_kind_lut: marks which adapter indices are classifier slots, so + # split_indices can route those positions out of the LoRA stream. None when + # no slot is a classifier. + kind_lut = build_adapter_kind_lut(config) + if kind_lut is not None: + self.register_buffer("adapter_kind_lut", kind_lut, persistent=False) + else: + self.adapter_kind_lut = None + lut = build_control_to_substitute_lut(config) if lut is not None: self.register_buffer("control_to_substitute_lut", lut, persistent=True) @@ -308,6 +322,15 @@ def _apply(self, *args, **kwargs): self.codebook = self.codebook.to(torch.float32) return out + def split_indices( + self, adapter_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Split ``[batch, seq_len]`` or ``[num_tokens]`` into (lora, classifier). + + See :func:`granite_switch.token_exchange.split_adapter_indices`. + """ + return split_adapter_indices(self.adapter_kind_lut, adapter_indices) + @property def num_cache_layers(self) -> int: """Cache slots used by this switch: counting slot + memory slot.""" diff --git a/src/granite_switch/token_exchange.py b/src/granite_switch/token_exchange.py index 231a2154..e0a37dff 100644 --- a/src/granite_switch/token_exchange.py +++ b/src/granite_switch/token_exchange.py @@ -154,3 +154,47 @@ def apply_token_exchange( sub_id_per_pos = lut[input_ids] is_control = sub_id_per_pos >= 0 return torch.where(is_control, sub_id_per_pos, input_ids) + + +def build_adapter_kind_lut(config) -> torch.Tensor | None: + """Bool LUT over adapter indices marking which slots are classifiers. + + ``[num_adapters + 1]``, indexed by ADAPTER INDEX: entry ``i`` is True when + adapter index ``i`` is a classifier slot (index 0, base, is always False). + ``None`` when no slot is a classifier, which lets the switch skip the split + entirely. Built from ``config.adapter_kinds``, so both backends' switches + agree on the table by construction. + """ + kinds = getattr(config, "adapter_kinds", None) if config is not None else None + if kinds is None or not any(k == "classifier" for k in kinds): + return None + lut = torch.zeros(len(kinds) + 1, dtype=torch.bool) + for slot, kind in enumerate(kinds): # slot is the 0-based adapter position + if kind == "classifier": + lut[slot + 1] = True # adapter index is slot + 1 + return lut + + +def split_adapter_indices( + kind_lut: torch.Tensor | None, adapter_indices: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Split a per-token index stream into LoRA and classifier streams. + + Shapes are preserved (``[batch, seq_len]`` on HF, ``[total_tokens]`` on + vLLM), so this serves both backends: + + - ``lora_indices``: classifier positions zeroed, so the LoRA path no-ops + there. + - ``classifier_indices``: the adapter index at classifier positions, else 0. + + With ``kind_lut`` None (no classifier slots) the LoRA stream passes through + and the classifier stream is all zeros. + """ + if kind_lut is None: + return adapter_indices, torch.zeros_like(adapter_indices) + is_classifier = kind_lut[adapter_indices] + zeros = torch.zeros_like(adapter_indices) + return ( + torch.where(is_classifier, zeros, adapter_indices), + torch.where(is_classifier, adapter_indices, zeros), + ) diff --git a/src/granite_switch/vllm/core/__init__.py b/src/granite_switch/vllm/core/__init__.py index 96afa952..894e7eb6 100644 --- a/src/granite_switch/vllm/core/__init__.py +++ b/src/granite_switch/vllm/core/__init__.py @@ -8,11 +8,13 @@ - lora_kernel_meta: Bitmask + remap metadata for the fused kernel """ +from .classifier import SwitchedClassifierHead from .lora import SwitchedLoRALinear from .lora_kernel_meta import FusedLoRAKernelMeta, LoRAContext __all__ = [ "FusedLoRAKernelMeta", "LoRAContext", + "SwitchedClassifierHead", "SwitchedLoRALinear", ] diff --git a/src/granite_switch/vllm/core/classifier.py b/src/granite_switch/vllm/core/classifier.py new file mode 100644 index 00000000..bf021ad9 --- /dev/null +++ b/src/granite_switch/vllm/core/classifier.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Switched classifier heads for Granite Switch (vLLM backend). + +A classifier head is an optional substitute for a LoRA adapter: an adapter slot +is either a LoRA adapter or a classifier head, selected by the same control +token and adapter index. Where a classifier is selected, the LoRA path no-ops +and the classifier reads clean base-model hidden states. + +Classifier positions are expected to be sparse, and the weight bank is replicated +on every rank at this point; revisit if classifiers ever cover most tokens or +num_labels grows large. +""" + +import torch +import torch.nn as nn + + +class SwitchedClassifierHead(nn.Module): + """Per-token classifier heads selected by classifier indices (vLLM). + + Compile-safe, on the flat ``[total_tokens, hidden_size]`` layout vLLM's + scheduler produces. Index ``0`` = no classifier (base -> zero logits); + ``1+`` selects that slot's head. + + The weight bank is stacked ``[num_slots, ...]`` and indexed by + ``(index - 1)``, matching ``SwitchedLoRALinear``. + + Args: + hidden_size: Input feature dimension (the model hidden size). + num_classifier_slots: ``num_adapters`` (the full index space). Slot ``s`` + (0-based) serves adapter index ``s + 1``. Slots that are not + classifiers simply never get selected. + max_num_labels: Padded label width of the bank (``max`` over the per-slot + label counts). Keeping this uniform is what lets the per-token logits + be one fused ``einsum``. Slots with fewer labels leave their extra + rows zero. The verdict exit slices each slot back to its real label + count. + """ + + def __init__( + self, + hidden_size: int, + num_classifier_slots: int, + max_num_labels: int, + ): + super().__init__() + self.hidden_size = hidden_size + self.num_classifier_slots = num_classifier_slots + self.num_labels = max_num_labels + + # Stacked bank indexed by (adapter_idx - 1), like the LoRA bank + # weight: [num_slots, max_num_labels, hidden_size] + # bias: [num_slots, max_num_labels] + # + # A frozen inference bank whose real values are copied in from a trained + # head at compose time. Zero value for both an unloaded slot's row and a + # loaded slot's pad rows (label counts below max_num_labels). + self.weight = nn.Parameter( + torch.zeros(num_classifier_slots, max_num_labels, hidden_size) + ) + self.bias = nn.Parameter(torch.zeros(num_classifier_slots, max_num_labels)) + + def forward( + self, + x: torch.Tensor, + classifier_indices: torch.Tensor, + ) -> torch.Tensor: + """Compute per-token classifier logits for selected positions. + + Branchless and loop-free so it traces cleanly under + ``@support_torch_compile``. + + Args: + x: Input tensor [total_tokens, hidden_size]. + classifier_indices: [total_tokens]; 0 = no classifier, 1+ = slot + index (the un-offset adapter index). + + Returns: + Logits [total_tokens, num_labels]. Positions with index 0 are zero. + """ + is_classifier = classifier_indices > 0 # [total_tokens] + slot = (classifier_indices - 1).clamp_min(0) # [total_tokens]; base -> row 0 + + w_per_token = self.weight[slot] # [total_tokens, num_labels, hidden] + b_per_token = self.bias[slot] # [total_tokens, num_labels] + + # Per-token linear, batched over tokens as one fused einsum. + logits = torch.einsum("tlh,th->tl", w_per_token, x) + logits = logits + b_per_token + + # Zero out base positions. + logits = logits * is_classifier.unsqueeze(-1).to(logits.dtype) + return logits diff --git a/src/granite_switch/vllm/core/lora_kernel_meta.py b/src/granite_switch/vllm/core/lora_kernel_meta.py index f7b740b9..94f6ca6f 100644 --- a/src/granite_switch/vllm/core/lora_kernel_meta.py +++ b/src/granite_switch/vllm/core/lora_kernel_meta.py @@ -192,6 +192,7 @@ class LoRAContext: __slots__ = ( "adapter_indices", + "classifier_indices", "num_tokens", "per_module_bitmasks", "remapped_indices", @@ -202,12 +203,18 @@ def __init__(self): self.remapped_indices: torch.Tensor | None = None self.per_module_bitmasks: torch.Tensor | None = None self.num_tokens: int = 0 + # Classifier stream (parallel to the LoRA stream): 0 = no classifier, + # 1+ = classifier slot. Set by the model's forward when a classifier + # head is present; read by the verdict exit on the last rank. Not part + # of the fused-LoRA bitmask path — the LoRA kernel never sees it. + self.classifier_indices: torch.Tensor | None = None def reset(self): self.adapter_indices = None self.remapped_indices = None self.per_module_bitmasks = None self.num_tokens = 0 + self.classifier_indices = None class FusedLoRAKernelMeta(nn.Module): diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index a6dd1c5b..e7f4450a 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -168,6 +168,23 @@ def __init__( self.lora_meta = None self.lora_ctx = None + # Classifier heads: optional alternative for LoRA adapters. A classifier + # slot shares the adapter index/token machinery, but its index is tagged + # as "classifier" (via the switch's adapter_kind_lut). The switch's + # split_indices() routes those positions out of the LoRA stream (so LoRA + # no-ops there) and into a classifier stream, which this head consumes + # from the final post-norm hidden state. + if "classifier" in getattr(config, "adapter_kinds", []): + from .core import SwitchedClassifierHead + + self.classifier_head = SwitchedClassifierHead( + hidden_size=config.hidden_size, + num_classifier_slots=config.num_adapters, + max_num_labels=config.max_classifier_labels, + ) + else: + self.classifier_head = None + # 3. Base transformer layers with custom LoRA # # When adapters are present, config.num_hidden_layers includes placeholder @@ -317,13 +334,30 @@ def forward( ) modified_input_ids = input_ids - # Prepare kernel metadata ONCE for all decoder layers (adaptation-specific). + # Split the single adapter_indices stream into a LoRA stream and a + # classifier stream. Classifier positions are zeroed in lora_indices + # (so the LoRA path no-ops there); the classifier head reads them on + # the last rank via classifier_indices. + if self.switch is not None: + lora_indices, classifier_indices = self.switch.split_indices( + adapter_indices + ) + else: + lora_indices = adapter_indices + classifier_indices = torch.zeros_like(adapter_indices) + + # Prepare kernel metadata ONCE for all decoder layers + # (adaptation-specific). The LoRA stream (classifier positions zeroed) + # is what the fused kernel routes on, so classifier tokens no-op the + # LoRA path. if self.lora_meta is not None and self.lora_ctx is not None: self.decoder_interface.prepare_kernel_meta( - self.lora_meta, adapter_indices, self.lora_ctx + self.lora_meta, lora_indices, self.lora_ctx ) # Store metadata in intermediate_tensors for pipeline parallelism. + # Only adapter_indices is propagated; classifier_indices is derived + # from it on the last rank via split_indices (no new PP tensor). if intermediate_tensors is None: intermediate_tensors = IntermediateTensors({}) intermediate_tensors["adapter_indices"] = adapter_indices @@ -332,9 +366,21 @@ def forward( # token-leading adapter_indices received through PP. if intermediate_tensors is not None: adapter_indices = intermediate_tensors["adapter_indices"] + # Recompute the split from the PP-propagated stream so this + # rank's Punica metadata sees classifier positions as base. + if self.switch is not None: + lora_indices, classifier_indices = self.switch.split_indices( + adapter_indices + ) + else: + lora_indices = adapter_indices + classifier_indices = torch.zeros_like(adapter_indices) + if self.lora_ctx is not None: + # Route on the LoRA stream (classifier positions zeroed) so + # classifier tokens no-op the LoRA path on this PP rank too. self.decoder_interface.prepare_kernel_meta( - self.lora_meta, adapter_indices, self.lora_ctx + self.lora_meta, lora_indices, self.lora_ctx ) else: # Fallback: no metadata available (should not happen in normal operation) @@ -350,6 +396,7 @@ def forward( dtype=torch.long, device=fallback_device, ) + classifier_indices = torch.zeros_like(adapter_indices) # ═══════════════════════════════════════════════════════════════ # Get embeddings (or hidden states from previous pipeline stage) @@ -384,10 +431,17 @@ def forward( if get_pp_group().is_last_rank: # Adaptation-specific finalize: LoRA folds the last residual via # rms_norm_select (fused/separate convention); SR does the per-token - # where-merge of its two streams then a plain norm. - return self.decoder_interface.exit_decoder_stack( + # where-merge of its two streams then a plain norm. The returned + # [M, H] tensor is the final post-norm hidden state (what the LM head + # consumes). + hidden_states = self.decoder_interface.exit_decoder_stack( state, adapter_indices, self.norm, self.config ) + + if self.lora_ctx is not None: + self.lora_ctx.classifier_indices = classifier_indices + + return hidden_states else: # Non-last rank: ship two token-leading [M,H] tensors so vLLM's # per-token IntermediateTensors slice stays correct. LoRA sends @@ -399,6 +453,145 @@ def forward( intermediate_tensors["residual"] = r_out return intermediate_tensors + def _classifier_read_points( + self, + classifier_indices: torch.Tensor, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Per-request (hidden_state, slot_index) for the classifier verdict. + + The verdict is read at the LAST CONTENT TOKEN, which the classifier + control token (marker) locates: the chat template emits the marker + immediately after the last content token, so the read point is + ``marker - 1``. + + The marker is the LAST classifier control token in a request's slice, + matched by id in ``input_ids``; ``classifier_indices`` is nonzero from any + classifier marker onward, so the index alone does not say which slot a + position belongs to. The read point is ``marker - 1`` and the slot id is + read AT the marker (``marker - 1`` carries index 0). + + Per request slice ``[start, end)``, three states: + + * marker present with a predecessor in this pass -- the common case. + * marker opens the slice: its predecessor ended the previous pass, so each + pass stashes its final hidden row for a single-token look-back. + * no marker: either still mid-prefill, or the marker was in an earlier + chunk and that pass's resolved verdict is carried forward, since vLLM + consumes the verdict only on the pass it samples. + + Returns ``None`` when there is no request metadata (e.g. the startup + profiling forward), so the caller emits no verdict. + """ + from vllm.forward_context import get_forward_context + + attn_metadata = get_forward_context().attn_metadata + # v1 keys this by layer name (a dict), or a list of such dicts under + # microbatching. They all share query_start_loc, so grab any one. + if isinstance(attn_metadata, list): + attn_metadata = attn_metadata[0] if attn_metadata else None + if isinstance(attn_metadata, dict): + attn_metadata = next(iter(attn_metadata.values()), None) + query_start_loc = getattr(attn_metadata, "query_start_loc", None) + if query_start_loc is None: + return None + + qsl = query_start_loc.tolist() + seq_lens = attn_metadata.seq_lens.tolist() + + # A request's marker is the LAST classifier control token in its slice: + # ``classifier_indices`` is nonzero from any classifier marker onward, so + # with several classifier slots the index alone does not say which slot's + # marker a position belongs to. + if input_ids is None: + raise RuntimeError( + "A classifier request needs input_ids to locate its control token, " + "but only inputs_embeds was provided. Classifier slots classify a " + "text prompt; pass input_ids." + ) + marker_ids = torch.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + is_marker = torch.isin(input_ids, marker_ids) # [total_tokens] + + stash = getattr(self, "_classifier_prev_pass_tail", None) or {} + resolved = getattr(self, "_classifier_resolved_verdict", None) or {} + req_hidden_rows = [] + req_slots = [] + new_stash: dict[int, torch.Tensor] = {} + new_resolved: dict[int, tuple[torch.Tensor, int]] = {} + for i in range(len(qsl) - 1): + start, end = qsl[i], qsl[i + 1] + hits = is_marker[start:end].nonzero(as_tuple=True)[0] + if hits.numel() == 0: + # No marker in this slice, which is either of two states: + # + # * the marker is still ahead (mid-prefill): keep this slice's final + # row, keyed by its global end position, in case the marker opens + # the next chunk. + # * the marker was in an earlier chunk (a rendered prompt continues + # past it with the generation prompt): carry that pass's resolved + # (row, slot) forward. vLLM consumes the verdict only on the pass + # where it samples, which is the final chunk. + carried = resolved.get(seq_lens[i] - (end - start)) + if carried is not None: + row, slot = carried + new_resolved[seq_lens[i]] = (row, slot) + req_hidden_rows.append(row) + req_slots.append(slot) + continue + new_stash[seq_lens[i]] = hidden_states[end - 1] + req_hidden_rows.append(hidden_states[end - 1]) + req_slots.append(0) + continue + + # One verdict per request: the exit rewrites a single logit row, so a + # request naming two different classifier slots has no way to report + # both. Repeated markers of the SAME slot are fine (a multi-turn prompt + # carries them from earlier turns); the last one wins. + distinct = torch.unique(input_ids[start:end][hits]) + if distinct.numel() > 1: + raise RuntimeError( + "A classifier request may name only one classifier slot; " + f"request {i} carries control tokens " + f"{sorted(int(t) for t in distinct)}. The verdict exit rewrites " + "one logit row, so only one slot can report." + ) + + marker = start + int(hits[-1]) # last id match: this slice's marker + slot = int(classifier_indices[marker]) + req_slots.append(slot) + if marker > start: + # Common case: predecessor is in this pass. + row = hidden_states[marker - 1] + new_resolved[seq_lens[i]] = (row, slot) + req_hidden_rows.append(row) + continue + + # Edge case: marker is the first token of its slice + marker_global = seq_lens[i] - (end - start) + prev = stash.get(marker_global) + if prev is None: + raise RuntimeError( + "Classifier verdict cannot locate its read point: the marker " + f"is the first token of request {i}'s chunk (global position " + f"{marker_global}), so the last content token (marker - 1) was " + "computed in the previous pass, and no stashed hidden row is " + "available for it." + ) + new_resolved[seq_lens[i]] = (prev, slot) + req_hidden_rows.append(prev) + + self._classifier_prev_pass_tail = new_stash + self._classifier_resolved_verdict = new_resolved + return ( + torch.stack(req_hidden_rows, dim=0), + torch.tensor(req_slots, dtype=torch.long, device=hidden_states.device), + ) + @MULTIMODAL_REGISTRY.register_processor( GraniteSwitchASRMultiModalProcessor, @@ -589,6 +782,23 @@ def forward( intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, ) + self._pending_verdict = None + classifier_head = getattr(self.model, "classifier_head", None) + if classifier_head is not None and isinstance(hidden_states, torch.Tensor): + classifier_indices = getattr( + self.model.lora_ctx, "classifier_indices", None + ) + if classifier_indices is not None: + read = self.model._classifier_read_points( + classifier_indices, hidden_states, input_ids + ) + if read is not None: + req_hidden, req_classifier_indices = read + verdict = classifier_head( + req_hidden, req_classifier_indices + ) # [num_reqs, num_labels] + self._pending_verdict = (verdict, req_classifier_indices) + return hidden_states def compute_logits( @@ -605,8 +815,51 @@ def compute_logits( per-token adapter_indices computed in forward(). Any suppression would have to act where adapter_indices and hidden_states still share the same token dimension. + + Classifier exit: the classifier request reports its verdict as a generated + token, like a guardian LoRA. For each classifier request we rewrite its logit + row: set the whole vocab to -inf, then place the per-label scores on each + label word's token id (``classifier_label_token_ids``) so the sampler + emits that word. Non-classifier rows keep their normal LM logits. """ - return self.logits_processor(self.lm_head, hidden_states) + logits = self.logits_processor(self.lm_head, hidden_states) + return self._apply_classifier_verdict(logits) + + def _apply_classifier_verdict( + self, + logits: torch.Tensor | None, + ) -> torch.Tensor | None: + """Rewrite classifier requests' logit rows to emit their label word.""" + pending = getattr(self, "_pending_verdict", None) + if logits is None or pending is None: + return logits + # Consume it so a same-step prompt_logprobs pass re-applies nothing. + self._pending_verdict = None + verdict, req_indices = pending + + per_slot_label_ids = self.config.classifier_label_token_ids + + if verdict.shape[0] != logits.shape[0]: + raise RuntimeError( + f"classifier verdict rows ({verdict.shape[0]}) do not match " + f"logits rows ({logits.shape[0]}); per-request alignment broke." + ) + + is_classifier = req_indices > 0 # [num_reqs] + if not bool(is_classifier.any()): + return logits + + # Each request writes only its own slot's label ids and only that slot's + # real label count (verdict[:, :n]); padded columns are never read. + rows = is_classifier.nonzero(as_tuple=True)[0] # [num_classifier_reqs] + logits[rows] = float("-inf") + for r in rows.tolist(): + slot = int(req_indices[r]) + ids = per_slot_label_ids[slot - 1] + n = len(ids) + label_ids = torch.tensor(ids, dtype=torch.long, device=logits.device) + logits[r, label_ids] = verdict[r, :n].to(logits.dtype) + return logits def sample( self, diff --git a/src/granite_switch/vllm/switch/multi.py b/src/granite_switch/vllm/switch/multi.py index 56a6d828..a4606b4e 100644 --- a/src/granite_switch/vllm/switch/multi.py +++ b/src/granite_switch/vllm/switch/multi.py @@ -104,7 +104,9 @@ ) from granite_switch.token_exchange import ( apply_token_exchange, + build_adapter_kind_lut, build_control_to_substitute_lut, + split_adapter_indices, ) # Large negative finite value for masking. NOT literal -inf because IEEE 754 @@ -312,12 +314,30 @@ def __init__( # non-persistent buffer is zeroed by checkpoint loading. This LUT uses -1 # as the "not a control token" sentinel, so an all-zero LUT would rewrite # EVERY token id to 0 in apply_token_exchange. + # adapter_kind_lut: marks which adapter indices are classifier slots, so + # split_indices can route those positions out of the LoRA stream. None when + # no slot is a classifier. + kind_lut = build_adapter_kind_lut(hf_config) + if kind_lut is not None: + self.register_buffer("adapter_kind_lut", kind_lut, persistent=False) + else: + self.adapter_kind_lut = None + lut = build_control_to_substitute_lut(hf_config) if lut is not None: self.register_buffer("control_to_substitute_lut", lut, persistent=True) else: self.control_to_substitute_lut = None + def split_indices( + self, adapter_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Split ``[total_tokens]`` into (lora_indices, classifier_indices). + + See :func:`granite_switch.token_exchange.split_adapter_indices`. + """ + return split_adapter_indices(self.adapter_kind_lut, adapter_indices) + @property def num_cache_layers(self) -> int: """KV-cache slots used by this switch: counting slot + memory slot.""" diff --git a/tests/composer/test_adapter_loader.py b/tests/composer/test_adapter_loader.py index 4d633a6b..372506d0 100644 --- a/tests/composer/test_adapter_loader.py +++ b/tests/composer/test_adapter_loader.py @@ -537,7 +537,9 @@ def test_fallback_precedence_and_yaml_parity(self, tmp_path): yaml.dump(manifest_data, f) input_path = str(manifest_file) - yaml_adapters = discover_adapters_from_yaml(input_path) + yaml_adapters, classifier_meta = discover_adapters_from_yaml(input_path) + # No entry declared kind/labels, so no classifier metadata. + assert classifier_meta == {} # Verification of Parity: compare (path, name, technology) - source differs by design # discover_adapters returns source=None, discover_adapters_from_yaml returns source=manifest_path @@ -547,3 +549,41 @@ def test_fallback_precedence_and_yaml_parity(self, tmp_path): adapters_without_source = [(p, n, t) for p, n, t, _ in adapters] yaml_without_source = [(p, n, t) for p, n, t, _ in yaml_adapters] assert sorted(adapters_without_source) == sorted(yaml_without_source) + + +class TestManifestClassifierMetadata: + """discover_adapters_from_yaml surfaces optional classifier declarations.""" + + def test_kind_and_labels_are_parsed(self, tmp_path): + manifest = tmp_path / "manifest.yaml" + yaml.dump( + { + "rag": {"path": "/p/rag", "type": "lora"}, + "hallucination": { + "path": "/p/hall", + "type": "lora", + "kind": "classifier", + "labels": ["faithful", "hallucinated"], + }, + }, + manifest.open("w"), + ) + + found, classifier_meta = discover_adapters_from_yaml(str(manifest)) + + # The adapter tuples are unchanged (still 4-wide). + assert {(n) for _, n, _, _ in found} == {"rag", "hallucination"} + # Only the classifier entry appears in the metadata dict. + assert classifier_meta == { + "hallucination": { + "kind": "classifier", + "labels": ["faithful", "hallucinated"], + } + } + + def test_plain_manifest_has_no_classifier_metadata(self, tmp_path): + manifest = tmp_path / "manifest.yaml" + yaml.dump({"rag": {"path": "/p/rag", "type": "lora"}}, manifest.open("w")) + + _found, classifier_meta = discover_adapters_from_yaml(str(manifest)) + assert classifier_meta == {} diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py new file mode 100644 index 00000000..a5eff866 --- /dev/null +++ b/tests/composer/test_classifier_compose.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compose-time loading of trained classifier-head weights. + +A "classifier" slot is the classifier-head analog of a LoRA adapter: it shares +the same control-token / adapter-index machinery, but instead of LoRA weights it +carries a small linear head (``weight [n_labels, hidden]`` + ``bias [n_labels]``) +that emits a per-request verdict. Each classifier slot may declare a DIFFERENT +number of labels; the head bank is padded to ``max_classifier_labels`` (like the +LoRA bank pads ranks to ``max_lora_rank``), so a slot's trained head lands in the +top ``n_labels`` rows of its bank row and the pad rows stay zero. The compose +pipeline builds the ``SwitchedClassifierHead`` from config and transfers the +trained head weights from each classifier directory's +``classifier_head.safetensors`` into bank row ``slot = adapter_index - 1``. + +These tests build the model through :class:`GraniteSwitchComposer` (per CLAUDE.md +rule #5 — never hand-assemble the config) against the real ``granite-4.0-micro`` +base, with a *synthetic* classifier head whose weights make one label win +deterministically, and assert: + +1. ``classifier_head.weight[slot, :n_labels]`` / ``.bias[slot, :n_labels]`` exactly + equal the synthetic tensors, and the pad rows above ``n_labels`` stay zero. +2. A classifier control-token forward emits the winning label's token id (the HF + verdict exit scatters the head's argmax onto that slot's + ``classifier_label_token_ids``). +3. In a mixed ``[lora, classifier]`` compose the LoRA slot's bank row is populated + while the classifier slot's LoRA row stays zero (compose orders LoRAs first, so + slot numbering is stable), and the classifier row lands in its slot. +4. Two classifier slots with DIFFERENT label counts (2 and 3) each land padded + correctly and each emit a token id from their own label set. + +Marked slow + requires_model: composing the ~3B micro base is expensive, but it is +the only sanctioned construction path. +""" + +import glob +import json +import os + +import pytest +import torch + +import granite_switch.hf # noqa: F401 — registers AutoModel/AutoConfig +from granite_switch.composer.compose_utils import GraniteSwitchComposer +from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE + +BASE_MODEL = "ibm-granite/granite-4.0-micro" +# The mixed test embeds a real LoRA, which must match its base. The RAG library +# ships granite-4.1-3b LoRAs, so the mixed case composes against that base. +MIXED_BASE = "ibm-granite/granite-4.1-3b" +MIXED_LORA_TARGET = "granite-4.1-3b" +RAG_LIBRARY = "ibm-granite/granite-lib-rag-r1.0" + +pytestmark = [pytest.mark.slow, pytest.mark.requires_model] + + +# --------------------------------------------------------------------------- # +# Helpers +# --------------------------------------------------------------------------- # +def _tokenizer(model=BASE_MODEL): + from transformers import AutoTokenizer + + try: + return AutoTokenizer.from_pretrained(model) + except Exception as e: # pragma: no cover - network/env dependent + pytest.skip(f"could not fetch tokenizer {model!r}: {e}") + + +def _hidden_size(model=BASE_MODEL): + from transformers import AutoConfig + + try: + return AutoConfig.from_pretrained(model).hidden_size + except Exception as e: # pragma: no cover + pytest.skip(f"could not fetch config {model!r}: {e}") + + +def _single_token_labels(tokenizer, candidates, count=2): + """Pick ``count`` label words that each encode to exactly one token id. + + The classifier verdict exit writes at one logit position per request, so a + label must be a single token. Returns ``(names, ids)``. + """ + names, ids = [], [] + for word in candidates: + enc = tokenizer.encode(word, add_special_tokens=False) + if len(enc) == 1 and enc[0] not in ids: + names.append(word) + ids.append(enc[0]) + if len(names) == count: + return names, ids + pytest.skip(f"could not find {count} single-token labels among {candidates}") + + +def _write_synthetic_head(directory, weight, bias, labels): + """Write a classifier slot dir: classifier_head.safetensors (+ config).""" + from safetensors.torch import save_file + + os.makedirs(directory, exist_ok=True) + save_file( + {"weight": weight.contiguous(), "bias": bias.contiguous()}, + os.path.join(directory, CLASSIFIER_HEAD_FILE), + ) + # Optional manifest metadata — the composer reads kinds/labels from kwargs, + # not this file, but a real classifier dir carries it, so include it. + with open(os.path.join(directory, "adapter_config.json"), "w") as f: + json.dump({"kind": "classifier", "labels": list(labels)}, f) + + +def _winning_head(num_labels, hidden, winner, dtype=torch.float32): + """Head whose winning label wins for ANY hidden state (bias-only). + + All weight rows are zero, so each label's logit reduces to its bias alone — + independent of the (real, arbitrary-magnitude) hidden state. Row ``winner`` + gets a large positive bias and the rest a large negative one, so ``argmax`` + over labels is deterministically ``winner``. A zero weight row is a valid + trained artifact shape and is the only construction that is robust to the + hidden state's sign and scale on a real base model. + """ + weight = torch.zeros((num_labels, hidden), dtype=dtype) + bias = torch.full((num_labels,), -10.0, dtype=dtype) + bias[winner] = 10.0 + return weight, bias + + +def _classifier_prompt_ids(tokenizer, control_token_id): + """`` as a batch of input ids for a detect forward. + + End-locator layout, matching what the chat template emits: the control token + (marker) goes AFTER the last content token, so the verdict is read at + ``marker - 1`` -- the position the head is trained on. A marker at position 0 + has no content token before it and is rejected by the backend. + """ + body = tokenizer.encode("Is this text safe?", add_special_tokens=False) + ids = [*body, control_token_id] + return torch.tensor([ids], dtype=torch.long) + + +def _find_rag_lora_dir(): + """Locate a real cached LoRA adapter dir for the mixed-case test. + + Reads the local HF cache directly (via ``scan_cache_dir``) rather than + ``snapshot_download``: the RAG library is large and only partially cached + here, and ``snapshot_download(local_files_only=True)`` refuses an incomplete + snapshot even when the one adapter this test needs is present. Skip cleanly + if no matching adapter is cached. + """ + from huggingface_hub import scan_cache_dir + + snapshot_dirs = [] + try: + for repo in scan_cache_dir().repos: + if repo.repo_id == RAG_LIBRARY: + snapshot_dirs.extend(str(rev.snapshot_path) for rev in repo.revisions) + except Exception as e: # pragma: no cover + pytest.skip(f"could not scan HF cache for {RAG_LIBRARY!r}: {e}") + + for root in snapshot_dirs: + hits = glob.glob( + os.path.join( + root, "**", MIXED_LORA_TARGET, "lora", "adapter_model.safetensors" + ), + recursive=True, + ) + # Follow the cache symlink to the real blob; skip broken (missing) ones. + hits = [h for h in sorted(hits) if os.path.exists(os.path.realpath(h))] + if hits: + return os.path.dirname(hits[0]) + + pytest.skip( + f"no cached {MIXED_LORA_TARGET} LoRA adapter with weights in {RAG_LIBRARY!r}" + ) + + +# --------------------------------------------------------------------------- # +# Tests +# --------------------------------------------------------------------------- # +def test_classifier_only_compose_loads_trained_head(tmp_path): + """A single classifier slot: weights land in the bank and the verdict fires.""" + tokenizer = _tokenizer() + hidden = _hidden_size() + + label_names, label_ids = _single_token_labels( + tokenizer, ["safe", "unsafe", "yes", "no", "true", "false"] + ) + num_labels = len(label_names) + winner = 1 # second label wins + weight, bias = _winning_head(num_labels, hidden, winner) + + clf_dir = str(tmp_path / "detect") + _write_synthetic_head(clf_dir, weight, bias, label_names) + + # A control token id for the single classifier slot. Any spare vocab id + # works for construction; the switch fires on this id, which the prompt + # helper places after the content (end-locator layout). + control_id = tokenizer.vocab_size - 5 + + model = GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=BASE_MODEL, + adapter_paths=[clf_dir], + adapter_token_ids=[control_id], + adapter_substitute_token_ids=[control_id], # inert for a classifier slot + adapter_names=["detect"], + adapter_kinds=["classifier"], + # Full-width per-slot label count (single classifier slot -> [num_labels]); + # ragged label metadata is one list per classifier slot. + classifier_label_token_ids=[label_ids], + ) + model.eval() + + # (1) The trained head is in bank row slot = index - 1 = 0, top n_labels rows. + head = model.model.classifier_head + assert head is not None, "classifier_head was not built" + torch.testing.assert_close( + head.weight[0, :num_labels].float(), weight.to(head.weight.dtype).float() + ) + torch.testing.assert_close( + head.bias[0, :num_labels].float(), bias.to(head.bias.dtype).float() + ) + + # (2) A detect forward emits the winning label's token id. + input_ids = _classifier_prompt_ids(tokenizer, control_id).to( + next(model.parameters()).device + ) + with torch.no_grad(): + out = model(input_ids=input_ids, use_cache=False) + emitted = int(out.logits[0, -1].argmax().item()) + assert emitted == label_ids[winner], ( + f"expected winning label id {label_ids[winner]} ({label_names[winner]}), " + f"got {emitted}" + ) + + +def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): + """[lora, classifier]: LoRA slot 0 populated, classifier slot 1 loaded.""" + tokenizer = _tokenizer(MIXED_BASE) + hidden = _hidden_size(MIXED_BASE) + lora_dir = _find_rag_lora_dir() + + label_names, label_ids = _single_token_labels( + tokenizer, ["safe", "unsafe", "yes", "no", "true", "false"] + ) + num_labels = len(label_names) + winner = 0 + weight, bias = _winning_head(num_labels, hidden, winner) + + clf_dir = str(tmp_path / "detect") + _write_synthetic_head(clf_dir, weight, bias, label_names) + + lora_ctrl = tokenizer.vocab_size - 6 + clf_ctrl = tokenizer.vocab_size - 5 + + model = GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=MIXED_BASE, + adapter_paths=[lora_dir, clf_dir], + adapter_token_ids=[lora_ctrl, clf_ctrl], + adapter_substitute_token_ids=[lora_ctrl, clf_ctrl], + adapter_names=["rag_lora", "detect"], + adapter_kinds=["lora", "classifier"], + # Full width: None on the LoRA slot, ids on the classifier slot. + classifier_label_token_ids=[None, label_ids], + ) + model.eval() + + head = model.model.classifier_head + assert head is not None + + # Classifier slot is index 1 -> bank row 1; slot 0 (the LoRA slot) is never + # touched by the classifier transfer, so it stays at the bank's zero init. + torch.testing.assert_close( + head.weight[1, :num_labels].float(), weight.to(head.weight.dtype).float() + ) + torch.testing.assert_close( + head.bias[1, :num_labels].float(), bias.to(head.bias.dtype).float() + ) + # Slot 0's classifier row must stay all-zero (never loaded) — a stronger check + # than "not equal to the synthetic head" now that the bank is zero-init. + assert head.weight[0].abs().sum() == 0, ( + "the LoRA slot's classifier row was overwritten" + ) + assert head.bias[0].abs().sum() == 0, ( + "the LoRA slot's classifier bias was overwritten" + ) + + # The LoRA bank row for slot 0 must be populated, while the classifier + # slot's LoRA row (slot 1) must stay unloaded. lora_B is the discriminator: + # it is ZERO-initialized and only becomes nonzero when an adapter is loaded + # (lora_A is kaiming-initialized, so an unloaded row is nonzero there too and + # can't distinguish loaded from unloaded). So: slot 0 lora_B nonzero (loaded), + # slot 1 lora_B all-zero (classifier slot, never received LoRA weights). + lora_B = _first_lora_b_tensor(model) + assert lora_B is not None, "no stacked LoRA B tensor found on the model" + assert lora_B[0].abs().sum() > 0, "LoRA slot 0 bank row is unexpectedly all-zero" + assert lora_B[1].abs().sum() == 0, ( + "classifier slot 1 LoRA row is not zero — slot numbering leaked" + ) + + # A detect forward on the classifier control token emits the winning label. + input_ids = _classifier_prompt_ids(tokenizer, clf_ctrl).to( + next(model.parameters()).device + ) + with torch.no_grad(): + out = model(input_ids=input_ids, use_cache=False) + emitted = int(out.logits[0, -1].argmax().item()) + assert emitted == label_ids[winner] + + +def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): + """Two classifier slots with DIFFERENT label counts (2 and 3). + + Each head lands in the top rows of its (padded) bank row, the pad rows above + its own label count stay zero, and a detect forward on each slot's control + token emits a token id from THAT slot's label set — proving the padded + einsum/scatter never leaks a phantom (padded) label across slots. + """ + tokenizer = _tokenizer() + hidden = _hidden_size() + + # Slot A: 2 labels; slot B: 3 labels. Distinct token ids across both sets. + names_a, ids_a = _single_token_labels( + tokenizer, ["safe", "unsafe", "yes", "no", "true", "false"], count=2 + ) + names_b, ids_b = _single_token_labels( + tokenizer, + ["good", "bad", "neutral", "high", "low", "medium", "left", "right", "up"], + count=3, + ) + n_a, n_b = len(names_a), len(names_b) # 2, 3 + max_labels = max(n_a, n_b) + + winner_a, winner_b = 1, 2 # last label of each set wins + w_a, b_a = _winning_head(n_a, hidden, winner_a) + w_b, b_b = _winning_head(n_b, hidden, winner_b) + + dir_a = str(tmp_path / "detect_a") + dir_b = str(tmp_path / "detect_b") + _write_synthetic_head(dir_a, w_a, b_a, names_a) + _write_synthetic_head(dir_b, w_b, b_b, names_b) + + ctrl_a = tokenizer.vocab_size - 5 + ctrl_b = tokenizer.vocab_size - 6 + + model = GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=BASE_MODEL, + adapter_paths=[dir_a, dir_b], + adapter_token_ids=[ctrl_a, ctrl_b], + adapter_substitute_token_ids=[ctrl_a, ctrl_b], + adapter_names=["detect_a", "detect_b"], + adapter_kinds=["classifier", "classifier"], + classifier_label_token_ids=[ids_a, ids_b], + ) + model.eval() + + head = model.model.classifier_head + assert head is not None + # Bank is padded to max_labels; each slot's head fills its top n rows. + assert head.weight.shape[1] == max_labels + + torch.testing.assert_close( + head.weight[0, :n_a].float(), w_a.to(head.weight.dtype).float() + ) + torch.testing.assert_close( + head.bias[0, :n_a].float(), b_a.to(head.bias.dtype).float() + ) + torch.testing.assert_close( + head.weight[1, :n_b].float(), w_b.to(head.weight.dtype).float() + ) + torch.testing.assert_close( + head.bias[1, :n_b].float(), b_b.to(head.bias.dtype).float() + ) + + # Slot A has fewer labels than max: its pad rows (n_a..max-1) must stay zero. + assert torch.count_nonzero(head.weight[0, n_a:]) == 0, "slot A pad rows not zero" + assert torch.count_nonzero(head.bias[0, n_a:]) == 0, "slot A pad bias not zero" + + device = next(model.parameters()).device + # Slot A fires -> emits one of A's ids (never a B id, never a padded slot). + with torch.no_grad(): + out_a = model( + input_ids=_classifier_prompt_ids(tokenizer, ctrl_a).to(device), + use_cache=False, + ) + emitted_a = int(out_a.logits[0, -1].argmax().item()) + assert emitted_a == ids_a[winner_a], f"slot A emitted {emitted_a}, want {ids_a}" + assert emitted_a not in ids_b, "slot A leaked a slot-B label" + + # Slot B fires -> emits one of B's ids, including its 3rd (only-in-B) label. + with torch.no_grad(): + out_b = model( + input_ids=_classifier_prompt_ids(tokenizer, ctrl_b).to(device), + use_cache=False, + ) + emitted_b = int(out_b.logits[0, -1].argmax().item()) + assert emitted_b == ids_b[winner_b], f"slot B emitted {emitted_b}, want {ids_b}" + + +def _first_lora_b_tensor(model): + """Return the first stacked LoRA-B bank tensor [num_adapters, 1, out, r]. + + lora_B is zero-init and only nonzero once an adapter is loaded, so it cleanly + distinguishes a loaded LoRA slot from an unloaded (classifier) slot — unlike + lora_A, which is kaiming-init and nonzero even when never loaded. + """ + for name, buf in model.named_parameters(): + if "lora_B" in name and buf.dim() == 4: + return buf.detach() + for name, buf in model.named_buffers(): + if "lora_B" in name and buf.dim() == 4: + return buf.detach() + return None diff --git a/tests/composer/test_tokenizer_setup.py b/tests/composer/test_tokenizer_setup.py index 36ca607b..d072dd91 100644 --- a/tests/composer/test_tokenizer_setup.py +++ b/tests/composer/test_tokenizer_setup.py @@ -12,6 +12,7 @@ add_control_tokens, build_substitute_token_ids, configure_chat_template, + resolve_label_token_ids, ) _PATCH_TARGET = "granite_switch.composer.tokenizer_setup._decode_alora_invocation_text" @@ -601,3 +602,46 @@ def test_multi_token_anchor_raises(self, tmp_path): with pytest.raises(ValueError, match="encodes to 3 tokens"): configure_chat_template(tokenizer, [(path, "answerability", "sr")]) + + +class _LabelTokenizer: + """Minimal tokenizer stub for resolve_label_token_ids. + + ``single_token`` maps a word -> its single id; any word not present is + treated as multi-token (splits into one id per character) so the multi-token + failure path can be exercised deterministically. + """ + + def __init__(self, single_token: dict[str, int]): + self._single = single_token + + def encode(self, text, add_special_tokens=False): + if text in self._single: + return [self._single[text]] + # Deterministic multi-token split: one synthetic id per character. + return [1000 + i for i, _ in enumerate(text)] + + def convert_ids_to_tokens(self, ids): + return [f"" for i in ids] + + +class TestResolveLabelTokenIds: + """Tests for resolve_label_token_ids (classifier label word -> token id).""" + + def test_single_token_labels_resolve_in_order(self): + tok = _LabelTokenizer({"safe": 19193, "unsafe": 39257}) + assert resolve_label_token_ids(tok, ["safe", "unsafe"]) == [19193, 39257] + + def test_order_is_preserved(self): + tok = _LabelTokenizer({"yes": 9891, "no": 2201}) + assert resolve_label_token_ids(tok, ["no", "yes"]) == [2201, 9891] + + def test_multi_token_label_fails_loud(self): + tok = _LabelTokenizer({"safe": 19193}) # "unsafe" absent -> multi-token + with pytest.raises(ValueError, match="single token"): + resolve_label_token_ids(tok, ["safe", "unsafe"]) + + def test_error_names_the_offending_word(self): + tok = _LabelTokenizer({"safe": 1}) + with pytest.raises(ValueError, match="Ambiguous"): + resolve_label_token_ids(tok, ["Ambiguous"]) diff --git a/tests/hf/test_classifier.py b/tests/hf/test_classifier.py new file mode 100644 index 00000000..80fd58d4 --- /dev/null +++ b/tests/hf/test_classifier.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Classifier-substitute path tests for GraniteSwitchForCausalLM (HF backend). + +A classifier slot is an alternative to a LoRA adapter: it shares the adapter +index/control-token machinery, the switch splits classifier positions out of +the LoRA stream, and the classifier head reads the final hidden state to emit +ONE verdict per request. The control token is placed after the last content +token, and the verdict is read at ``marker - 1`` (the last content token). + +The verdict exits as a GENERATED LABEL WORD, matching the vLLM backend: the LM +head rewrites the classifier request's last logit row so the vocab is -inf +except at ``classifier_label_token_ids``, where the per-label verdict scores are +placed. There is no separate float ``classifier_logits`` output field. + +CPU, random weights, no checkpoint — mirrors tests/hf/test_model_forward.py. +""" + +import pytest +import torch + +from granite_switch.config import GraniteSwitchConfig +from granite_switch.hf import GraniteSwitchForCausalLM + + +def _set_adapter_token_ids(model, token_ids): + model.model.adapter_token_ids.data = torch.tensor(token_ids, dtype=torch.long) + + +def _classifier_config(adapter_kinds, num_labels=2, label_token_ids=(100, 200)): + # Full-width label token ids (length num_adapters, None on LoRA slots). Every + # classifier slot gets the same ids here; mixed counts are covered by the + # compose tests. + label_token_ids = list(label_token_ids) + return GraniteSwitchConfig( + vocab_size=300, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, # 1 switch + 2 decoder + num_attention_heads=4, + num_key_value_heads=4, + num_adapters=2, + adapter_token_ids=[250, 251], + adapter_substitute_token_ids=[1, 1], + adapter_names=["adapter_1", "adapter_2"], + adapter_kinds=adapter_kinds, + classifier_label_token_ids=[ + list(label_token_ids) if k == "classifier" else None for k in adapter_kinds + ], + max_lora_rank=4, + adapter_ranks=[4, 4], + switch_head_dim=16, + ) + + +def test_classifier_head_built_only_when_slot_present(): + m_lora = GraniteSwitchForCausalLM(_classifier_config(["lora", "lora"])) + m_clf = GraniteSwitchForCausalLM(_classifier_config(["lora", "classifier"])) + assert m_lora.model.classifier_head is None + assert m_clf.model.classifier_head is not None + # Bank sized to the full index space so slot = index - 1 addresses it. + assert m_clf.model.classifier_head.num_classifier_slots == 2 + assert m_clf.model.classifier_head.num_labels == 2 + + +def test_classifier_rewrites_last_logit_row_to_label_word(): + # adapter 2 (token 251) is a classifier. Its verdict rides the LM logits: + # the last prompt row is -inf everywhere except at the label token ids, + # where the per-label verdict scores land — so decoding emits a label word. + config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + + # Marker (251) at the end (end-locator layout); this test checks the label-word + # rewrite of the emitted row. The read position is pinned in + # test_classifier_reads_marker_minus_one_hidden_state. + input_ids = torch.tensor([[10, 20, 30, 40, 50, 60, 70, 251]]) + with torch.no_grad(): + out = model(input_ids=input_ids) + + last_row = out.logits[0, -1, :] # [vocab] + # Single classifier slot -> classifier position 0; its own label ids. + label_ids = config.classifier_label_token_ids[1] # slot 2 -> index 1 + n = len(label_ids) + non_label = [t for t in range(config.vocab_size) if t not in label_ids] + + # Every non-label token is blanked; the label ids carry the verdict. + assert torch.isinf(last_row[non_label]).all() + assert (last_row[non_label] < 0).all() + assert torch.isfinite(last_row[torch.tensor(label_ids)]).all() + # The label-id scores equal the classifier head's per-label verdict, sliced + # to this slot's real label count (padded columns are never scattered). + verdict = model.model._last_classifier_logits[0, :n] # [n_labels] + assert torch.allclose(last_row[torch.tensor(label_ids)], verdict.to(last_row.dtype)) + # The argmax (emitted token) is one of the label ids. + assert int(last_row.argmax()) in label_ids + + +def test_non_classifier_request_logits_untouched(): + # No classifier token (251) present -> last token is base (index 0), the + # head's last-token verdict row is zero, and NO row is blanked: a + # non-classifier request keeps a full finite vocab row. + config = _classifier_config(["lora", "classifier"]) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + + input_ids = torch.tensor([[10, 20, 30, 40]]) + with torch.no_grad(): + out = model(input_ids=input_ids) + + assert (model.model._last_classifier_indices == 0).all() + # Full finite vocab row — nothing blanked. + assert torch.isfinite(out.logits[0, -1, :]).all() + + +def test_mixed_batch_only_classifier_row_rewritten(): + # Batch: row 0 is a classifier request (251 fires), row 1 is plain LM. + config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + + input_ids = torch.tensor( + [ + [10, 20, 251, 40, 50], # classifier + [10, 20, 30, 40, 50], # plain + ] + ) + with torch.no_grad(): + out = model(input_ids=input_ids) + + label_ids = torch.tensor(config.classifier_label_token_ids[1]) # slot 2 + non_label = [ + t for t in range(config.vocab_size) if t not in set(label_ids.tolist()) + ] + + # Classifier row rewritten; plain row fully finite. + assert torch.isinf(out.logits[0, -1, non_label]).all() + assert torch.isfinite(out.logits[1, -1, :]).all() + + +def test_classifier_is_detect_only_decode_step_raises(): + # A classifier request is prefill-only: it classifies the prompt and stops. + # A cached decode step (generation) on a classifier model is misuse and + # must fail loud rather than produce a verdict over generated tokens. + config = _classifier_config(["lora", "classifier"]) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + + input_ids = torch.tensor([[10, 20, 251, 40, 50]]) + with torch.no_grad(): + out = model(input_ids=input_ids, use_cache=True) + assert out.logits.shape == (1, 5, config.vocab_size) + + with pytest.raises(RuntimeError, match="detect-only"): + with torch.no_grad(): + model( + input_ids=torch.tensor([[60]]), + past_key_values=out.past_key_values, + use_cache=True, + ) + + +def test_classifier_missing_label_token_ids_fails_loud(): + # A classifier slot with no resolved label token ids cannot emit a label word. + # Rejected when the config is built, so a bad checkpoint fails at load rather + # than at the first inference. + with pytest.raises(ValueError, match="classifier_label_token_ids"): + GraniteSwitchConfig( + vocab_size=300, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=4, + num_adapters=2, + adapter_token_ids=[250, 251], + adapter_substitute_token_ids=[1, 1], + adapter_names=["adapter_1", "adapter_2"], + adapter_kinds=["lora", "classifier"], + classifier_label_token_ids=None, + max_lora_rank=4, + adapter_ranks=[4, 4], + switch_head_dim=16, + ) + + +def test_lora_only_path_unaffected(): + # A model with no classifier slot behaves exactly as before: no classifier + # head, no verdict, full finite LM logits. + config = _classifier_config(["lora", "lora"]) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + + input_ids = torch.tensor([[10, 250, 20, 251, 30, 40]]) + with torch.no_grad(): + out = model(input_ids=input_ids) + + assert model.model.classifier_head is None + assert model.model._last_classifier_logits is None + assert torch.isfinite(out.logits).all() + assert out.logits.shape == (1, 6, config.vocab_size) + + +def test_lora_generation_on_mixed_model_decodes(): + # A model composed WITH a classifier slot also serves plain LoRA/base + # generation. Firing the LoRA (adapter 1, token 250) — not the classifier — + # must be allowed to decode multiple tokens: the detect-only guard keys on an + # actually-fired classifier, not on the mere presence of a classifier head + # bank. (Regression: the guard previously raised on ANY cached decode step of + # a classifier-composed model, crashing all multi-token LoRA generation.) + config = _classifier_config(["lora", "classifier"]) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + assert model.model.classifier_head is not None # mixed model has the bank + + # Prefill fires the LoRA (250), never the classifier (251). + input_ids = torch.tensor([[10, 250, 20, 30, 40]]) + with torch.no_grad(): + out = model(input_ids=input_ids, use_cache=True) + # No classifier fired -> no verdict, full finite vocab row. + assert model.model._last_classifier_logits is None + assert torch.isfinite(out.logits).all() + + # Cached decode step of the same (LoRA) request must NOT raise and must + # yield finite logits — classifier_indices stays zero, so the detect-only + # guard does not fire. + with torch.no_grad(): + step = model( + input_ids=torch.tensor([[60]]), + past_key_values=out.past_key_values, + use_cache=True, + ) + assert step.logits.shape == (1, 1, config.vocab_size) + assert torch.isfinite(step.logits).all() + assert model.model._last_classifier_logits is None + + +def test_classifier_reads_marker_minus_one_hidden_state(): + """The verdict is read at ``marker - 1`` (the last content token), not at the + marker itself or the sequence's last token. + + End-locator layout: the classifier control token is placed after the last content + token, so ``classifier_indices`` is nonzero only from the marker onward (the switch + forward-fills the index causally), and the read point is ``marker - 1``. The exact + position is pinned by capturing (a) the final post-norm hidden states via a hook on + ``model.norm`` and (b) the hidden state the classifier head actually receives, then + asserting the head's input equals the post-norm hidden at + ``marker - 1`` — and differs from the marker's and the last token's. + """ + config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + # Populate the classifier bank (slot 1 = adapter index 2) so the head is a real, + # non-zero map (a zero bank would make every position's verdict identical). + with torch.no_grad(): + model.model.classifier_head.weight[1].normal_(0.0, 0.3) + model.eval() + + # Marker (251) at the LAST position -> read point = marker - 1 = position 4. + input_ids = torch.tensor([[10, 20, 30, 40, 60, 251]]) + marker_pos = 5 + read_pos = marker_pos - 1 # 4, the last content token + + captured = {} + + # Capture the final post-norm hidden states (exactly what the classifier reads). + def norm_hook(_module, _inp, output): + captured["normed"] = output.detach().clone() + + h = model.model.norm.register_forward_hook(norm_hook) + + # Spy on the head input so we see which hidden state it was handed. + real_head = model.model.classifier_head + orig_forward = real_head.forward + + def spy_forward(x, classifier_indices): + captured["head_x"] = x.detach().clone() + captured["head_slot"] = classifier_indices.detach().clone() + return orig_forward(x, classifier_indices) + + real_head.forward = spy_forward + try: + with torch.no_grad(): + model(input_ids=input_ids) + finally: + real_head.forward = orig_forward + h.remove() + + normed = captured["normed"][0] # [seq_len, hidden] + head_x = captured["head_x"] # [batch, hidden] (gathered read-point state) + assert head_x.shape == (1, config.hidden_size) + + # The head was handed the post-norm hidden at marker - 1, not marker or last token. + assert torch.allclose(head_x[0], normed[read_pos], atol=1e-5), ( + "classifier head did not read the last content token (marker - 1)" + ) + assert not torch.allclose(head_x[0], normed[marker_pos], atol=1e-5), ( + "classifier head read the marker itself, not marker - 1" + ) + # marker_pos is also the sequence's last position, so this pins that the read + # is NOT the last token. + assert read_pos != marker_pos + + # The slot id is read AT the marker (marker - 1 carries index 0). + assert int(captured["head_slot"][0]) == 2, ( + "classifier slot id must be read at the marker (adapter index 2)" + ) + + +def test_two_classifier_slots_in_one_request_rejected(): + """A request may name only one classifier slot. + + The verdict exit rewrites a single logit row, so two slots cannot both report. + ``classifier_indices`` is nonzero from either marker onward, so silently + taking one would compute the verdict from whichever position won the race. + """ + config = _classifier_config( + ["classifier", "classifier"], label_token_ids=(100, 200) + ) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + model.eval() + + first_tok, last_tok = config.adapter_token_ids + with pytest.raises(RuntimeError, match="only one classifier slot"): + with torch.no_grad(): + model(input_ids=torch.tensor([[first_tok, 10, 20, 30, last_tok]])) + + +def test_repeated_marker_for_one_slot_reads_the_last(): + """The same slot's marker may repeat; the last occurrence is the read point. + + A multi-turn prompt carries markers from earlier turns, so a repeat is + legitimate and must not be rejected -- only the current turn's marker decides + where the verdict is read. + """ + config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + model = GraniteSwitchForCausalLM(config) + _set_adapter_token_ids(model, config.adapter_token_ids) + with torch.no_grad(): + model.model.classifier_head.weight[1].normal_(0.0, 0.3) + model.eval() + + cls_tok = config.adapter_token_ids[1] + input_ids = torch.tensor([[cls_tok, 10, 20, 30, cls_tok]]) + marker_pos = 4 + + captured = {} + real_head = model.model.classifier_head + orig_forward = real_head.forward + + def spy(x, classifier_indices): + captured["x"] = x.detach().clone() + return orig_forward(x, classifier_indices) + + real_head.forward = spy + hook = model.model.norm.register_forward_hook( + lambda _m, _i, out: captured.__setitem__("normed", out.detach().clone()) + ) + try: + with torch.no_grad(): + model(input_ids=input_ids) + finally: + real_head.forward = orig_forward + hook.remove() + + normed = captured["normed"][0] + assert torch.allclose(captured["x"][0], normed[marker_pos - 1], atol=1e-5), ( + "verdict was not read at the last marker's predecessor" + ) diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py new file mode 100644 index 00000000..ddb62e4b --- /dev/null +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Subprocess worker for the vLLM classifier chunked-prefill test. + +Each sub-command is a separate process invocation so only one vLLM engine is ever +resident (same discipline as ``_generation_equivalence_worker``):: + + python worker.py compose # CPU only; prints {"result": meta} + python worker.py compose-multilabel # CPU only; multi-label head + python worker.py run # one engine, one budget + python worker.py fault-sweep # forced-read sweep over the prompt + +The engine core runs in-process (the driver sets ``VLLM_ENABLE_V1_MULTIPROCESSING=0``) +so the hooks installed here land on the model the engine actually runs. + +The read-index hook records, per forward pass, the global position the verdict read. +``_classifier_read_points`` reads at ``marker - 1`` (a flat within-pass index); the hook +maps that index through ``positions``, vLLM's own per-token coordinate, so the recorded +position reflects the token actually read. +""" + +import json +import os +import sys +import traceback + +_REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +if _REPO not in sys.path: + sys.path.insert(0, _REPO) + + +def _install_read_index_hook(): + """Record, per forward pass, the global position the classifier verdict actually + read. + + Wraps ``GraniteSwitchModel._classifier_read_points`` to capture, per request, the + flat within-pass index it read (``marker - 1``, or a sentinel ``-1`` when the read + came from the cross-pass stash because the marker sat at the slice start). The + ``forward`` wrapper maps that index through ``positions`` to a global position. + """ + from granite_switch.vllm.granite_switch_model import ( + GraniteSwitchForCausalLM, + GraniteSwitchModel, + ) + + passes = [] + orig_fwd = GraniteSwitchForCausalLM.forward + orig_read = GraniteSwitchModel._classifier_read_points + + # Per-pass scratch for request 0 (these tests use a single request): + # flat_read_idx -- within-pass flat index the verdict read (marker-1), or None. + # stash_global -- when the read came from the previous-pass stash (marker at + # the slice start), the global position of that stashed token + # (== num_computed - 1, with num_computed derived as + # seq_lens - query_len); this pass's ``positions`` cannot show + # it, so it is carried here directly. + scratch = {"flat_read_idx": None, "stash_global": None} + + def read_points(self, classifier_indices, hidden_states, input_ids): + # Snapshot before the call: it replaces _classifier_resolved_verdict with + # this pass's map, and the carry check needs the previous pass's. + carried_before = dict(getattr(self, "_classifier_resolved_verdict", None) or {}) + res = orig_read(self, classifier_indices, hidden_states, input_ids) + try: + from vllm.forward_context import get_forward_context + + am = get_forward_context().attn_metadata + if isinstance(am, list): + am = am[0] if am else None + if isinstance(am, dict): + am = next(iter(am.values()), None) + qsl = getattr(am, "query_start_loc", None) + if qsl is not None: + qsl = qsl.tolist() + start, end = qsl[0], qsl[1] # request 0 + # The marker is the last control-token id match in this slice. + import torch as _t + + mids = _t.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + hits = _t.isin(input_ids[start:end], mids).nonzero(as_tuple=True)[0] + if hits.numel() == 0: + # No marker in this slice. If a verdict was resolved in an + # earlier pass it is carried forward, and the position it was + # read at is not addressable by this pass's ``positions``; report + # it directly. Otherwise this request is still mid-prefill. + sql = getattr(am, "seq_lens", None) + pre = ( + int(sql.tolist()[0]) - (end - start) + if sql is not None + else None + ) + if pre is not None and pre in carried_before: + scratch["flat_read_idx"] = -1 + scratch["stash_global"] = scratch.get("carried_global") + else: + scratch["flat_read_idx"] = end - 1 + else: + marker = start + int(hits[-1]) + if marker > start: + # Common case: predecessor (marker-1) is in this pass. + scratch["flat_read_idx"] = marker - 1 + scratch["resolving_flat"] = marker - 1 + else: + # Edge case: marker at slice start; the read came from the + # previous-pass stash. Its global position is + # num_computed - 1, where num_computed (tokens computed + # BEFORE this pass) is seq_lens - query_len. + sql = getattr(am, "seq_lens", None) + nc0 = ( + int(sql.tolist()[0]) - (end - start) + if sql is not None + else None + ) + scratch["flat_read_idx"] = -1 + scratch["stash_global"] = nc0 - 1 if nc0 is not None else None + except Exception: # a hook must never break a forward + scratch["flat_read_idx"] = None + return res + + def fwd(self, input_ids, positions, *args, **kwargs): + scratch["flat_read_idx"] = None + scratch["stash_global"] = None + out = orig_fwd(self, input_ids, positions, *args, **kwargs) + try: + has_cls = getattr(self.model, "classifier_head", None) is not None + flat = scratch["flat_read_idx"] + if has_cls and positions is not None and flat is not None: + if flat == -1: + read_pos = scratch["stash_global"] # from previous pass + else: + read_pos = int(positions[flat].item()) + if scratch.get("resolving_flat") is not None: + # Remember where the verdict was resolved, so a later pass that + # merely carries it can report the same global position. + scratch["carried_global"] = int( + positions[scratch["resolving_flat"]].item() + ) + scratch["resolving_flat"] = None + passes.append( + {"read_global_pos": read_pos, "n_query": int(positions.shape[0])} + ) + except Exception: # a hook must never break a forward + pass + return out + + GraniteSwitchModel._classifier_read_points = read_points + GraniteSwitchForCausalLM.forward = fwd + return passes + + +# Single-token label candidates, in priority order (only single-token ones are kept). +_LABEL_CANDIDATES = [ + "safe", + "unsafe", + "yes", + "no", + "true", + "false", + "good", + "bad", + "high", + "low", + "hot", + "cold", +] + + +def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): + """Compose a single-classifier-slot checkpoint. + + ``num_labels`` single-token labels are taken from ``_LABEL_CANDIDATES`` and + ``scale`` is the std of the random head weight. Both are set high enough, over a + non-repetitive prompt, that the argmax varies across read positions -- which is + what makes a wrong read observable in the fault sweep. + """ + import torch + from safetensors.torch import save_file + from transformers import AutoConfig, AutoTokenizer + + from granite_switch.composer.compose_utils import GraniteSwitchComposer + from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE + + tok = AutoTokenizer.from_pretrained(base) + hidden = AutoConfig.from_pretrained(base).hidden_size + + labels, ids = [], [] + for c in _LABEL_CANDIDATES: + e = tok.encode(c, add_special_tokens=False) + if len(e) == 1 and e[0] not in ids: + labels.append(c) + ids.append(e[0]) + if len(labels) == num_labels: + break + if len(labels) < num_labels: + raise RuntimeError( + f"only {len(labels)} single-token labels available from " + f"{_LABEL_CANDIDATES}, need {num_labels}" + ) + nl = len(labels) + + g = torch.Generator().manual_seed(0) + # Weight-based (not bias-only) head so the verdict depends on the hidden state. + weight = torch.randn((nl, hidden), generator=g, dtype=torch.float32) * scale + bias = torch.zeros((nl,), dtype=torch.float32) + + slot = os.path.join(outdir, "detect") + os.makedirs(slot, exist_ok=True) + save_file( + {"weight": weight.contiguous(), "bias": bias.contiguous()}, + os.path.join(slot, CLASSIFIER_HEAD_FILE), + ) + json.dump( + {"kind": "classifier", "labels": labels}, + open(os.path.join(slot, "adapter_config.json"), "w"), + ) + open(os.path.join(slot, "io.yaml"), "w").write("name: detect\nmodel: ~\n") + + control_id = tok.vocab_size - control_offset + model = GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=base, + adapter_paths=[slot], + adapter_token_ids=[control_id], + adapter_substitute_token_ids=[control_id], + adapter_names=["detect"], + adapter_kinds=["classifier"], + classifier_label_token_ids=[ids], + ) + ckpt = os.path.join(outdir, "composed") + model.save_pretrained(ckpt) + tok.save_pretrained(ckpt) + + # Non-repetitive body so mid-prompt hidden states differ across positions. + words = ( + "machine translation quietly reshaped how distant villages argued " + "about rainfall while engineers debated whether copper cables or " + "glass fibre carried gossip faster than a startled horse could run " + "downhill past the old mill where children once traded marbles for " + "stories about comets volcanoes submarines and the strange arithmetic " + "of tides that neither king nor merchant ever fully trusted yet " + "everyone quoted at dinner as though the ocean kept a ledger of debts " + "owed to the moon each evening without fail or apparent complaint " + "although sailors insisted otherwise over cheap wine and louder songs" + ) + body = tok.encode(f"Is this text safe? {words}", add_special_tokens=False) + # End-locator layout, matching the chat template: the marker follows the last + # content token, and the verdict is read at marker - 1. + prompt_ids = [*body, control_id] + # Same layout, but continuing past the marker the way a rendered chat prompt + # does (turn close + generation prompt). The marker is then interior, so a + # chunk budget can close the prompt after it -- the pass that samples holds no + # marker at all and must consume the verdict resolved in an earlier pass. + tail = tok.encode( + "\n<|start_of_role|>assistant<|end_of_role|>", add_special_tokens=False + ) + prompt_ids_trailing = [*body, control_id, *tail] + return { + "ckpt": ckpt, + "labels": labels, + "control_id": control_id, + "label_token_ids": ids, + "prompt_ids": prompt_ids, + "prompt_ids_trailing": prompt_ids_trailing, + "n_tail": len(tail), + } + + +def run_budget(ckpt, prompt_ids, budget, out_path): + import torch # noqa: F401 + from vllm import LLM, SamplingParams + from vllm.inputs import TokensPrompt + + passes = _install_read_index_hook() + kw = dict( + model=ckpt, + enforce_eager=True, + gpu_memory_utilization=0.55, + max_model_len=4096, + enable_prefix_caching=False, + dtype="bfloat16", + ) + if budget is None: + kw["enable_chunked_prefill"] = False + else: + kw["max_num_batched_tokens"] = budget + kw["enable_chunked_prefill"] = True + llm = LLM(**kw) + sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=5) + o = llm.generate([TokensPrompt(prompt_token_ids=prompt_ids)], sp)[0].outputs[0] + + # The consumed verdict is the final pass's; single request => one read per pass. + last = passes[-1] if passes else None + global_idx = None + if last and last["read_global_pos"] is not None: + rp = last["read_global_pos"] + global_idx = int(rp[0]) if isinstance(rp, list) else int(rp) + + out = { + "budget": budget, + "token_id": int(o.token_ids[0]), + "text": o.text, + "n_passes": len(passes), + "global_read_idx": global_idx, + "n_prompt": len(prompt_ids), + } + json.dump(out, open(out_path, "w")) + + +def run_fault_sweep(ckpt, prompt_ids, out_path): + """Force the classifier verdict to read positions swept across the whole prompt + and record the resulting token at each. + + The correct read point is ``marker - 1`` (the last content token), located inside + ``_classifier_read_points``. To sweep wrong reads, that method is replaced with a + forcing version: it takes the real per-request slice but reads + ``hidden_states[start + forced[0]]`` and reports slot 1 (the classifier slot) so the + verdict still fires. ``forced[0] is None`` runs the real method. + + The prefill runs unchunked (a single pass over all ``n`` tokens), so the within-pass + flat index equals the global position and a forced ``start + i`` reaches global + position ``i`` for any ``0 <= i <= n-1`` in one resident engine. A mutable + ``forced[0]`` cell, closed over by the monkeypatch, changes the forced index between + ``generate`` calls without rebuilding the engine.""" + import torch + from vllm import LLM, SamplingParams + from vllm.forward_context import get_forward_context + from vllm.inputs import TokensPrompt + + from granite_switch.vllm.granite_switch_model import GraniteSwitchModel + + n = len(prompt_ids) + + forced = [None] # None => real method (marker-1); int => forced within-slice index + read_global = [None] # global position the forced read landed on + orig_read = GraniteSwitchModel._classifier_read_points + + def patched_read(self, classifier_indices, hidden_states, input_ids): + if forced[0] is None: + return orig_read(self, classifier_indices, hidden_states, input_ids) + am = get_forward_context().attn_metadata + if isinstance(am, list): + am = am[0] if am else None + if isinstance(am, dict): + am = next(iter(am.values()), None) + qsl = getattr(am, "query_start_loc", None) + if qsl is None: + return None + qsl = qsl.tolist() + start = qsl[0] # single request in these tests + flat = start + int(forced[0]) + # Unchunked single pass => flat within-pass index == global position. + read_global[0] = int(forced[0]) + req_hidden = hidden_states[flat : flat + 1] # [1, hidden] + # Slot 1 = the classifier slot, so the verdict fires like a real read. + req_slot = torch.ones(1, dtype=torch.long, device=hidden_states.device) + return req_hidden, req_slot + + GraniteSwitchModel._classifier_read_points = patched_read + try: + passes = _install_read_index_hook() # wraps the (now patched) method + forward + llm = LLM( + model=ckpt, + enforce_eager=True, + gpu_memory_utilization=0.55, + max_model_len=4096, + enable_prefix_caching=False, + dtype="bfloat16", + enable_chunked_prefill=False, + ) + sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=5) + pt = TokensPrompt(prompt_token_ids=prompt_ids) + + def one(forced_idx): + forced[0] = forced_idx + read_global[0] = None + passes.clear() + o = llm.generate([pt], sp)[0].outputs[0] + if forced_idx is None: + # Correct run: read position comes from the hook (marker - 1). + last = passes[-1] if passes else None + rp = last["read_global_pos"] if last else None + pos = int(rp) if rp is not None else None + else: + # Forced run: the global position is the forced index directly. + pos = read_global[0] + return {"pos": pos, "token_id": int(o.token_ids[0]), "text": o.text} + + # Correct run: reads marker - 1 (global pos n-2, marker is at n-1). + correct = one(None) + + # Faulted sweep at a fixed stride over the prompt, excluding the correct + # read point (n-2) so every forced read is genuinely a wrong position. + step = max(1, (n - 1) // 10) + sweep_idx = sorted({i for i in range(0, n - 1, step)} - {n - 2}) + faulted = [one(i) for i in sweep_idx] + finally: + GraniteSwitchModel._classifier_read_points = orig_read + + out = {"correct": correct, "faulted": faulted, "n_prompt": n} + json.dump(out, open(out_path, "w")) + + +def main(): + mode = sys.argv[1] + if mode == "compose-multilabel": + try: + meta = compose(os.environ["CLS_BASE"], os.environ["CLS_OUT"]) + print(json.dumps({"result": meta})) + except Exception as e: + print(json.dumps({"error": f"{e}\n{traceback.format_exc()}"})) + return + if mode == "fault-sweep": + out_path = sys.argv[2] + prompt_ids = json.load(open(os.environ["CLS_PIDS"])) + try: + run_fault_sweep(os.environ["CLS_CKPT"], prompt_ids, out_path) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return + if mode == "run": + budget = None if sys.argv[2] == "none" else int(sys.argv[2]) + out_path = sys.argv[3] + prompt_ids = json.load(open(os.environ["CLS_PIDS"])) + try: + run_budget(os.environ["CLS_CKPT"], prompt_ids, budget, out_path) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return + raise SystemExit(f"unknown mode {mode!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_classifier_chunked_prefill.py new file mode 100644 index 00000000..ab15a194 --- /dev/null +++ b/tests/vllm/test_classifier_chunked_prefill.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Chunked-prefill coverage for the vLLM classifier verdict exit. + +Every other vLLM classifier path is HF-only and single-pass, yet +``enable_chunked_prefill`` is on by default in serving. Two tests here: + +The prompt places the classifier control token immediately after the last +content token (marker at position ``n-1``), and the verdict is read at +``marker - 1`` (the last content token, ``n-2``) inside +``GraniteSwitchModel._classifier_read_points``. + +1. ``test_classifier_token_stable_across_chunk_budgets`` — the same prompt served at + ``max_num_batched_tokens`` in {unchunked, 512, 256, 128, 101, 64, 48, n_prompt-1} + emits the same verdict token, and each budget's consumed verdict is read from + ``marker - 1``. ``_classifier_read_points`` locates the marker per pass and reads + its predecessor; when a chunk boundary splits the marker from its predecessor, a + single-token cross-pass stash supplies the previous pass's last hidden row. The + ``n_prompt - 1`` budget is the arm that forces that split (marker alone in the + final chunk) — the only budget here that exercises the look-back at all. + +2. ``test_wrong_read_corrupts_position_and_often_verdict`` — the inverse: forcing the + read to wrong positions lands at a wrong global position (asserted) and usually + flips the label (reported; see that test). + +Not asserted: logprob or hidden-state equality across budgets. Chunked and unchunked +serving differ by bf16 reduction order in the attention/KV path (a magnitude that is +hardware- and version-specific), so a numeric band would be flaky. Token stability +still catches a wrong read position or a weight/logit corruption. + +Each budget runs in its own subprocess so only one vLLM engine is resident, with the +engine core in-process (``VLLM_ENABLE_V1_MULTIPROCESSING=0``) so the read-index hook +lands on the running model. + +Requires GPU + vLLM. +""" + +import json +import os +import subprocess +import sys + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="vLLM classifier serving test requires a GPU" +) + +WORKER = os.path.join( + os.path.dirname(__file__), "_classifier_chunked_prefill_worker.py" +) +_REPO = os.path.dirname(os.path.dirname(os.path.dirname(WORKER))) +BASE_MODEL = os.environ.get("CLS_TEST_BASE", "ibm-granite/granite-4.1-3b") + +# One budget per distinct final-chunk shape, which is what the read point depends on: +# ``None`` disables chunked prefill (single pass); 256 leaves a wide final chunk; 48 +# splits into many passes with a narrow one; 101 leaves exactly two tokens, so the +# marker and its read point are the last pair in the chunk. The lone-marker budget +# (final chunk of one) is derived per prompt and appended by the test below. Budgets +# differing only in pass count exercise no additional path. +BUDGETS = [None, 256, 101, 48] + + +def _lone_marker_budget(n_prompt: int) -> int: + """Budget that puts the marker ALONE in the final chunk. + + The marker is the prompt's last token, so it is the first (and only) token of + the final chunk exactly when ``n_prompt % budget == 1`` -- i.e. budget + ``n_prompt - 1``, which splits the prompt into [0, n-2] then [n-1]. That is + the one layout where the read point (``marker - 1``) lands in the *previous* + pass, exercising the cross-pass look-back in + ``GraniteSwitchModel._classifier_read_points``. Computed from ``n_prompt`` so a + tokenizer change that shifts it cannot silently stop testing this case. + """ + return n_prompt - 1 + + +def _env(): + env = dict(os.environ) + env["VLLM_USE_V1"] = "1" + env["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" # in-process engine core + env["TOKENIZERS_PARALLELISM"] = "false" + env.setdefault("PYTHONPATH", _REPO) + return env + + +def _compose(tmp_path_factory, mode, tag): + """Compose a classifier checkpoint in a CPU subprocess, before any CUDA init.""" + outdir = str(tmp_path_factory.mktemp(tag)) + proc = subprocess.run( + [sys.executable, WORKER, mode], + env={**_env(), "CLS_BASE": BASE_MODEL, "CLS_OUT": outdir}, + capture_output=True, + text=True, + timeout=1800, + ) + lines = [l for l in (proc.stdout or "").splitlines() if l.startswith("{")] + if proc.returncode != 0 or not lines: + pytest.fail( + f"{mode} failed:\nstdout={proc.stdout[-2000:]}\n" + f"stderr={proc.stderr[-2000:]}", + pytrace=False, + ) + meta = json.loads(lines[-1]) + if "error" in meta: + pytest.fail(f"{mode} error: {meta['error'][:2000]}", pytrace=False) + result = meta["result"] + result["pids_path"] = os.path.join(outdir, "pids.json") + json.dump(result["prompt_ids"], open(result["pids_path"], "w")) + result["pids_trailing_path"] = os.path.join(outdir, "pids_trailing.json") + json.dump(result["prompt_ids_trailing"], open(result["pids_trailing_path"], "w")) + result["outdir"] = outdir + return result + + +@pytest.fixture(scope="module") +def composed(tmp_path_factory): + """Multi-label classifier over a varied prompt: the argmax is position-sensitive, + which both tests need (stable-across-budgets and the forced-read sweep).""" + return _compose(tmp_path_factory, "compose-multilabel", "cls_ckpt") + + +def _run_worker(composed, argv, out_name, label, timeout, pids_key="pids_path"): + """Run one worker invocation in its own subprocess; return the parsed result.""" + out_path = os.path.join(composed["outdir"], out_name) + proc = subprocess.run( + [sys.executable, WORKER, *argv, out_path], + env={ + **_env(), + "CLS_CKPT": composed["ckpt"], + "CLS_PIDS": composed[pids_key], + }, + capture_output=True, + text=True, + timeout=timeout, + ) + if not os.path.exists(out_path): + pytest.fail( + f"{label}: worker wrote no result\n" + f"stdout={proc.stdout[-1500:]}\nstderr={proc.stderr[-1500:]}", + pytrace=False, + ) + res = json.load(open(out_path)) + if "fatal" in res: + pytest.fail(f"{label}: worker crashed:\n{res['fatal'][:2000]}", pytrace=False) + return res + + +def _run(composed, budget, pids_key="pids_path"): + """Serve the prompt at one chunk budget (``None`` = chunked prefill disabled).""" + tag = "none" if budget is None else str(budget) + suffix = "" if pids_key == "pids_path" else "_trailing" + return _run_worker( + composed, + ["run", tag], + f"run_{tag}{suffix}.json", + f"budget={budget}{suffix}", + 1200, + pids_key=pids_key, + ) + + +def _run_fault_sweep(composed): + """Run the full-prompt forced-read sweep in one subprocess.""" + return _run_worker( + composed, ["fault-sweep"], "fault_sweep.json", "fault-sweep", 1800 + ) + + +def test_classifier_token_stable_across_chunk_budgets(composed): + """Emitted verdict token identical across chunk budgets; every budget's + consumed verdict is read from ``marker - 1`` (the last content token). + + The prompt is ``[*body, control_id]`` (marker at ``n_prompt - 1``), so the + correct read point is ``marker - 1 == n_prompt - 2``. + + Budgets include ``n_prompt - 1``, the one budget that leaves the marker ALONE + in the final chunk (``n_prompt % budget == 1``). There the read point sits in + the *previous* pass, so this arm is what exercises the cross-pass look-back; + every other budget keeps marker and read point in the same chunk.""" + n_prompt = len(composed["prompt_ids"]) + read_point = n_prompt - 2 # marker is at n_prompt - 1; read point = marker - 1 + label_ids = set(composed["label_token_ids"]) + + # The lone-marker budget (n_prompt % budget == 1) is the only layout that + # splits the marker from its read point across a chunk boundary, so it is the + # one arm that exercises the cross-pass look-back. Derived from n_prompt. + lone = _lone_marker_budget(n_prompt) + budgets = [*BUDGETS, lone] + assert n_prompt % lone == 1, ( + f"budget {lone} does not leave the marker alone in the final chunk for " + f"n_prompt={n_prompt} (n_prompt % budget == {n_prompt % lone}, want 1)" + ) + + results = {b: _run(composed, b) for b in budgets} + + ref = results[None] + assert ref["token_id"] in label_ids, ( + f"unchunked verdict {ref['token_id']} is not a label token {label_ids}; " + f"the classifier exit did not fire" + ) + + # Consumed verdict reads the last content token (marker - 1), for every budget. + for b, r in results.items(): + assert r["global_read_idx"] == read_point, ( + f"budget={b}: consumed verdict read from global index " + f"{r['global_read_idx']}, expected {read_point} (marker-1) " + f"(passes={r['n_passes']})" + ) + + # Verdict token stable across budgets. + for b, r in results.items(): + assert r["token_id"] == ref["token_id"], ( + f"budget={b} emitted {r['token_id']} ({r['text']!r}) != unchunked " + f"{ref['token_id']} ({ref['text']!r})" + ) + + # The smallest budget actually chunked. + assert results[48]["n_passes"] > 1, ( + f"budget=48 ran in {results[48]['n_passes']} pass(es); expected chunking." + ) + + +def test_wrong_read_corrupts_position_and_often_verdict(composed): + """A wrong read position corrupts the verdict. + + Forces the verdict to read positions strided across the prompt and records the + emitted token at each, against a correct run reading ``marker - 1`` (the last + content token). + Asserts a wrong read always lands at a wrong global position (the correct run + reads ``n_prompt - 2``; every forced read reports a different position). + + The label-flip rate is reported, not asserted: a head can map two hidden states to + the same argmax label, so some wrong reads produce the correct label by chance, and + a mismatch threshold would be flaky. The multi-label head keeps the argmax + position-sensitive so most wrong reads do flip the label in practice. + """ + n_prompt = len(composed["prompt_ids"]) + read_point = n_prompt - 2 # marker at n_prompt - 1; correct read = marker - 1 + label_ids = set(composed["label_token_ids"]) + + sweep = _run_fault_sweep(composed) + correct = sweep["correct"] + faulted = sweep["faulted"] + + assert correct["pos"] == read_point, ( + f"correct run read global pos {correct['pos']}, expected last content token " + f"{read_point} (marker - 1)" + ) + assert correct["token_id"] in label_ids, ( + f"correct verdict {correct['token_id']} is not a label token {label_ids}; " + f"the classifier exit did not fire" + ) + # Expect a spread of distinct positions, every one inside 0..n-2 and none at + # the correct read point. + assert len(faulted) >= 3, ( + f"fault sweep produced only {len(faulted)} forced reads; expected a spread" + ) + positions_set = {f["pos"] for f in faulted} + assert all(0 <= p <= n_prompt - 2 for p in positions_set), ( + f"forced reads landed outside the prompt pool 0..{n_prompt - 2}: " + f"{sorted(positions_set)}" + ) + for f in faulted: + assert f["pos"] != read_point, ( + f"forced read reported global pos {f['pos']} == correct read point " + f"{read_point}; the forced index did not move the read" + ) + + # Report the label-flip rate (not asserted; see docstring). + n_mismatch = sum(1 for f in faulted if f["token_id"] != correct["token_id"]) + hist = {} + for f in faulted: + hist[f["text"]] = hist.get(f["text"], 0) + 1 + positions = sorted(f["pos"] for f in faulted) + frac = n_mismatch / len(faulted) if faulted else 0.0 + print( + f"\nfault-sweep: {n_mismatch}/{len(faulted)} wrong-read positions " + f"(global {positions[0]}..{positions[-1]}) produced a label != correct " + f"{correct['text']!r} ({frac:.0%}); labels={hist}" + ) + + +def test_verdict_carries_when_marker_is_not_in_the_sampling_chunk(composed): + """A marker in an earlier chunk still produces its verdict. + + A rendered chat prompt continues past the classifier marker (turn close plus + generation prompt), so the marker is interior. With a budget that closes a chunk + at the marker, the pass vLLM samples on contains no marker at all -- the verdict + was resolved one pass earlier and has to be carried forward. + + The prompt is ``[*body, marker, *tail]``, so the marker sits at ``n_tail`` tokens + from the end and its read point is ``marker - 1``. + """ + n_prompt = len(composed["prompt_ids_trailing"]) + n_tail = composed["n_tail"] + marker_pos = n_prompt - n_tail - 1 + read_point = marker_pos - 1 + label_ids = set(composed["label_token_ids"]) + assert n_tail > 0, "the trailing variant must continue past the marker" + + # Budget that ends a chunk exactly at the marker: the next pass starts at the + # first tail token and holds no marker. + budget = marker_pos + 1 + res = _run(composed, budget, pids_key="pids_trailing_path") + + assert res["n_passes"] > 1, ( + f"budget={budget} ran in {res['n_passes']} pass(es); expected the prompt to " + f"split after the marker" + ) + assert res["token_id"] in label_ids, ( + f"verdict {res['token_id']} is not a label token {label_ids}; the verdict " + f"resolved before the sampling chunk was not carried forward" + ) + assert res["global_read_idx"] == read_point, ( + f"carried verdict was read from global index {res['global_read_idx']}, " + f"expected {read_point} (marker at {marker_pos}, {n_tail} tail tokens)" + ) + + # Same prompt unchunked must agree. + ref = _run(composed, None, pids_key="pids_trailing_path") + assert res["token_id"] == ref["token_id"], ( + f"chunked verdict {res['token_id']} != unchunked {ref['token_id']} on the " + f"trailing-token prompt" + ) diff --git a/tutorials/README.md b/tutorials/README.md index 62d120c7..969b4541 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -26,6 +26,7 @@ Step-by-step walkthroughs covering adapter function invocation, pipeline constru |-------|-------------| | [Using Mellea with Granite Switch](guides/mellea_with_granite_switch.md) | Connect Mellea to a Granite Switch model | | [Bring Your Own Adapter](guides/build_your_own_adapter.md) | Train, compose, and use custom adapters | +| [Classifier Slots](guides/classifier_slots.md) | Compose and serve a classifier head that emits a label word | | [Compare Inference Throughput](guides/compare_inference_throughput.md) | Compare LoRA vs aLoRA based models in an inference race setup | diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md new file mode 100644 index 00000000..f39426a9 --- /dev/null +++ b/tutorials/guides/classifier_slots.md @@ -0,0 +1,297 @@ +# Classifier Slots + +This guide explains how to compose a Granite Switch model with a **classifier slot** — a +lightweight classification head fired by the same control token as a LoRA adapter — and how to +serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detection request. + +## Overview + +A Granite Switch adapter slot is normally a LoRA adapter. A **classifier slot** is an alternative +occupant of that same slot: instead of a LoRA, it holds a small linear classification head. It is +triggered by the *same* control-token machinery — one `<|name|>` token, one adapter index — but +where a LoRA would modify the hidden states, a classifier reads the clean base-model hidden state +and produces per-label logits. + +At inference the verdict does not come back as a raw score vector. Each label word is resolved to a +single token id at compose time, and the model rewrites the request's final logit row so that the +winning label's token id wins the argmax. The endpoint therefore emits the literal label word — the +same shape as any other generation, so it works through the standard vLLM OpenAI server with no +custom client. + +``` +1. Train a head ──> 2. Manifest (kind/labels) ──> 3. Compose ──> 4. Serve + detect + (weight/bias) (io.yaml or YAML list) (granite-switch) (vLLM OpenAI API) +``` + +Multiple classifier slots may coexist, and **each slot may have its own number of labels** (e.g. a +2-label `safe`/`unsafe` slot next to a 3-label `good`/`bad`/`neutral` slot). + +## Prerequisites + +- Granite base model (e.g., `ibm-granite/granite-4.1-3b`) +- A trained classifier head (see Step 1) +- `granite-switch[vllm,compose]` installed + +See [PREREQUISITES.md](../PREREQUISITES.md) for detailed setup. + +## Step 1: Train (or stage) a classifier head + +A classifier slot is **not** a LoRA, so its directory holds a different artifact. There is no +`adapter_model.safetensors`; instead the directory must contain a **`classifier_head.safetensors`** +with exactly two tensors: + +| Key | Shape | Meaning | +|-----|-------|---------| +| `weight` | `[num_labels, hidden_size]` | the linear head | +| `bias` | `[num_labels]` | per-label bias | + +`num_labels` is *this slot's own* label count, and `hidden_size` must match the base model's hidden +size. That is the whole artifact — a single linear layer that maps a base-model hidden state to one +logit per label. + +Train the head however you like (a frozen base model plus a trainable linear layer over the last +hidden state is the usual recipe). Here is a helper that writes a **synthetic** head. + +```python +from pathlib import Path +import torch +from safetensors.torch import save_file + +HIDDEN = 2560 # granite-4.1-3b hidden size + +def make_head(name, labels): + """Write a synthetic classifier head for `name` under the library layout. + + A real head would carry trained weights; this one is zero-weight with a bias + that forces row 0 to win, so the emitted label is deterministic. + """ + head_dir = Path(f"./{name}-head/{name}/granite-4.1-3b/lora") + head_dir.mkdir(parents=True, exist_ok=True) + weight = torch.zeros(len(labels), HIDDEN) + bias = torch.full((len(labels),), -10.0) + bias[0] = 10.0 # row 0 (labels[0]) always wins + save_file({"weight": weight, "bias": bias}, str(head_dir / "classifier_head.safetensors")) + # Every composed slot (LoRA or classifier) also needs an io.yaml — compose copies + # it verbatim into the checkpoint's io_configs/. A two-line stub is enough here. + (head_dir / "io.yaml").write_text(f"name: {name}\nmodel: ~\n") + return str(head_dir) +``` + +> The directory layout `////` mirrors the library layout the +> composer expects for a local adapter. The trailing `lora`/`alora` segment selects **token +> placement** (LoRA vs activated-LoRA), independent of the fact that this slot is a classifier. +> +> The `io.yaml` is required for **every** slot: compose copies each slot's `io.yaml` into the +> output's `io_configs//`, so a slot without one fails the copy step. Real adapters ship a +> full `io.yaml` (instruction, response schema, sampling parameters); a synthetic slot only needs +> the two-line stub above. + +### Label words must be single-token + +Each label word is resolved to a **single token id** at compose time — the verdict is written at one +token id per label. If a label word does not encode to a single token in the base model's tokenizer, +compose fails loud and lists the offenders, e.g.: + +``` +1 classifier label word(s) do not encode to a single token, ... + 'Ambiguous' → 3 tokens: ['Amb', 'igu', 'ous'] +``` + +Pick a single-token synonym for each label. Short lowercase words (`safe`, `unsafe`, `yes`, `no`, +`good`, `bad`, `neutral`) are single-token in Granite BPE; capitalized or rarer words often are not. +If no natural single-token word fits, the tokenizer's reserved `<|unused_N|>` tokens are a fallback. + +## Step 2: Declare the slot in a manifest + +Classifier slots are declared with two extra keys — `kind: classifier` and `labels:` — on top of the +normal adapter fields. The simplest way to pass one or more of them to compose is a **YAML manifest** +that maps each slot name to its path and metadata: + +```yaml +# ./adapters.yaml (one classifier slot) +safety: + path: ./safety-head/safety/granite-4.1-3b/lora + type: lora # token placement (lora vs alora); independent of `kind` + kind: classifier # marks this slot a classifier head (default "lora" when omitted) + labels: [safe, unsafe] # label words, resolved to single token ids at compose +``` + +`type` and `kind` are independent axes: `type` is the LoRA *technology* (drives control-token +placement), while `kind` is whether the slot is a real LoRA or a classifier head. A plain LoRA entry +simply omits `kind`/`labels`. + +You can mix LoRA and classifier entries in one manifest, and give each classifier its own `labels` +list of any length: + +```yaml +# ./adapters.yaml — one real LoRA + two synthetic classifier slots (2 and 3 labels) +answerability: + path: /path/to/answerability/granite-4.1-3b/lora # a real LoRA (see below) + type: lora +safety: + path: ./safety-head/safety/granite-4.1-3b/lora + type: lora + kind: classifier + labels: [safe, unsafe] +faithfulness: + path: ./faithfulness-head/faithfulness/granite-4.1-3b/lora + type: lora + kind: classifier + labels: [good, bad, neutral] # single-token words — see the label rule above +``` + +> **Slot ordering is automatic.** Compose stably reorders slots so all LoRA slots come first and +> classifier slots last, then derives every downstream artifact (control tokens, chat template, +> adapter indices) from that order. You do not need to order the manifest yourself; classifiers +> ending up as the contiguous suffix is what lets the LoRA weight transfer stay positional and the +> classifier head bank pad uniformly. + +The script below download a rag lora and constructed two synthetic classifiers for testing. + +```python +from pathlib import Path +import yaml +from huggingface_hub import snapshot_download +# make_head is the helper defined in Step 1. + +# 1. A real LoRA slot: download just the answerability adapter for granite-4.1-3b. +snap = snapshot_download( + repo_id="ibm-granite/granite-lib-rag-r1.0", + allow_patterns=["answerability/granite-4.1-3b/lora/**"], +) +lora_dir = str(Path(snap) / "answerability" / "granite-4.1-3b" / "lora") + +# 2. Two synthetic classifier slots with different label counts (make_head from Step 1). +safety_dir = make_head("safety", ["safe", "unsafe"]) +faith_dir = make_head("faithfulness", ["good", "bad", "neutral"]) + +# 3. Emit the manifest pairing the real LoRA with the two synthetic classifiers. +manifest = { + "answerability": {"path": lora_dir, "type": "lora"}, + "safety": {"path": safety_dir, "type": "lora", "kind": "classifier", + "labels": ["safe", "unsafe"]}, + "faithfulness": {"path": faith_dir, "type": "lora", "kind": "classifier", + "labels": ["good", "bad", "neutral"]}, +} +Path("./adapters.yaml").write_text(yaml.safe_dump(manifest, sort_keys=False)) +print("wrote ./adapters.yaml") +``` + + +## Step 3: Compose + +Pass the manifest to the compose CLI exactly like any other adapter source — a `.yaml`/`.yml` path is +auto-detected as a manifest: + +```bash +python -m granite_switch.composer.compose_granite_switch \ + --base-model ibm-granite/granite-4.1-3b \ + --adapters ./adapters.yaml \ + --output ./granite-switch-detect +``` + +Compose will: + +1. Load the base model and discover the slots from the manifest. +2. Resolve each classifier's label words to single token ids (failing loud on multi-token words). +3. Read each classifier's `classifier_head.safetensors` into the head bank at its slot. +4. Stack everything into one checkpoint under `--output`. + + +``` +Classifier slots: + safety → labels ['safe', 'unsafe'] → token ids [19193, 39257] + faithfulness → labels ['good', 'bad', 'neutral'] → token ids [19045, 14176, 60668] +``` + +The resulting `config.json` carries the per-slot classifier metadata +(`classifier_num_labels_per_slot`, `max_classifier_labels`, `classifier_label_names`, +`classifier_label_token_ids`, `classifier_slot_indices`) — all plain lists/ints, so it round-trips +through `save_pretrained`/`from_pretrained` with no custom code, and both the HF and vLLM backends +read it identically. + +## Step 4: Serve and send a detection request + +Serve the composed checkpoint with the standard vLLM OpenAI-compatible server: + +```bash +python -m vllm.entrypoints.openai.api_server \ + --model ./granite-switch-detect \ + --port 8000 +``` + +Fire a classifier by placing its control token in the prompt — one token per slot, `<|safety|>` for +the `safety` slot, `<|faithfulness|>` for the other. The completion is the emitted label word: + +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "./granite-switch-detect", + "prompt": "<|safety|>Is this text safe? The weather is nice today.", + "max_tokens": 1 + }' +``` + +```json +{"choices": [{"text": "safe", ...}], ...} +``` + +Only the classifier request's final logit row is rewritten; every other (non-classifier) request runs +untouched with its full vocabulary. If you request logprobs, you will see the label token ids +carrying the verdict scores and the rest of the vocabulary pinned at the blanking floor — exactly the +two labels for a 2-label slot, three for a 3-label slot. + +### Firing via the chat template (optional) + +Instead of placing `<|safety|>` by hand, you can let the composed model's chat template place it for +you by passing `adapter_name` — a classifier slot is activated exactly like any other adapter: + +```bash +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "./granite-switch-detect", + "messages": [{"role": "user", "content": "Is this text safe? The weather is nice today."}], + "max_tokens": 1, + "chat_template_kwargs": {"adapter_name": "safety"} + }' +``` + +### Reading logprobs (optional) + +To see the raw verdict rather than just the winning word, ask for a couple of logprobs: + +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "./granite-switch-detect", + "prompt": "<|safety|>Is this text safe? ...", + "max_tokens": 1, + "logprobs": 5 + }' +``` + +Only the label ids (`safe`, `unsafe`) will show finite logprobs; all other tokens sit at the blanking +floor. That is the classifier head's per-label verdict surfaced through the standard logprobs +channel. + +## How it works (brief) + +- **Same trigger as a LoRA.** One `<|name|>` control token per slot, one adapter index. The switch + tags that index as a classifier and routes the token out of the LoRA stream, so the LoRA path + no-ops there and the classifier head runs on the clean base hidden state. +- **Padded head bank.** All classifier heads live in one stacked tensor padded to + `max_classifier_labels`; a slot with fewer labels leaves its extra rows zero. This keeps the vLLM + head a single fused, compile-safe operation regardless of how many labels each slot has — the same + way the LoRA bank pads every adapter to `max_lora_rank`. +- **Label-word verdict exit.** The model collapses each classifier request to its final token, + computes the head's per-label logits, blanks the vocabulary, and scatters the verdict onto that + slot's own label token ids (sliced to the slot's real label count, so padded columns are never + emitted). The sampler then emits the label word. + +## Related + +- [Bring Your Own Adapter](build_your_own_adapter.md) — train and compose a LoRA adapter. +- [Composer README](../../src/granite_switch/composer/README.md) — full compose CLI reference. From ef335c049e31c9b24b7c52e6f9c2839c130ebd9e Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Mon, 28 Sep 2026 23:21:16 +0000 Subject: [PATCH 02/20] docs(classifier): list the config fields compose actually writes Update the Step 3 paragraph to match GraniteSwitchConfig: `adapter_kinds` and `classifier_label_token_ids` are the stored fields, with the per-slot label counts and control-token ids derived from them. Signed-off-by: IRENEKO --- tutorials/guides/classifier_slots.md | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index f39426a9..a4077dcc 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -204,11 +204,27 @@ Classifier slots: faithfulness → labels ['good', 'bad', 'neutral'] → token ids [19045, 14176, 60668] ``` -The resulting `config.json` carries the per-slot classifier metadata -(`classifier_num_labels_per_slot`, `max_classifier_labels`, `classifier_label_names`, -`classifier_label_token_ids`, `classifier_slot_indices`) — all plain lists/ints, so it round-trips -through `save_pretrained`/`from_pretrained` with no custom code, and both the HF and vLLM backends -read it identically. +The resulting `config.json` carries the per-slot classifier metadata as plain lists/ints, so it +round-trips through `save_pretrained`/`from_pretrained` with no custom code, and both the HF and +vLLM backends read it identically. Two fields are the stored contract: + +| Field | Meaning | +|-------|---------| +| `adapter_kinds` | per-slot `"lora"` / `"classifier"`, one entry per adapter — what marks a slot a classifier | +| `classifier_label_token_ids` | per-slot list of label token ids, `null` on LoRA slots | + +The rest are **derived** from those in `GraniteSwitchConfig.__init__`, so they appear on the config +object but are not independent inputs: + +| Derived field | From | +|---------------|------| +| `classifier_num_labels_per_slot` | `len()` of each slot's label token ids (`0` on LoRA slots) | +| `max_classifier_labels` | `max()` of the above — the padded width the head bank is built with | +| `classifier_control_token_ids` | the control-token ids of the classifier slots; both backends match these in `input_ids` to locate a verdict's read point | + +Note that the label *words* are not stored — only their token ids. Compose resolves and prints the +words, but `config.json` keeps ids, so recovering the words means decoding them with the +checkpoint's tokenizer. ## Step 4: Serve and send a detection request From ef553d4eba10268b35c211bd7975a3c5c8d0fe16 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 30 Sep 2026 16:08:51 +0000 Subject: [PATCH 03/20] refactor(classifier): read the verdict at the marker The read point was ``marker - 1``; it is now the marker's own position. The marker's row is always in the pass that contains the marker, so the cross-pass stash that served the one-token look-back is gone, along with its bookkeeping and the error for a missing stash entry. Signed-off-by: IRENEKO --- .../composer/tokenizer_setup.py | 7 +- .../hf/modeling_granite_switch.py | 17 +--- .../vllm/granite_switch_model.py | 59 +++---------- tests/composer/test_classifier_compose.py | 4 +- tests/hf/test_classifier.py | 68 ++++++++------ .../_classifier_chunked_prefill_worker.py | 58 ++++-------- tests/vllm/test_classifier_chunked_prefill.py | 88 ++++++------------- 7 files changed, 107 insertions(+), 194 deletions(-) diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 6e9cac57..455fbf31 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -1178,8 +1178,7 @@ def configure_chat_template( """ # Classifier placement: the marker goes after the last user turn's content and - # before that turn's closing marker (``<|end_of_text|>`` / ``<|im_end|>``), so - # the read point ``marker - 1`` is the position the head was trained on. + # before that turn's closing marker (``<|end_of_text|>`` / ``<|im_end|>``). classifier_scan = ( """{#- Classifier scan: find the last user message (the turn being classified). -#} {%- if ns.adapter_token and ns.adapter_type == 'classifier' %} @@ -1501,7 +1500,9 @@ def _split_merged(match: "re.Match") -> str: print(f"Chat template configured with {len(adapter_mapping)} adapter mappings:") for adapter_name, info in adapter_mapping.items(): if info["type"] == ANCHOR_MODE_CLASSIFIER: - placement = "at the end of the last user turn (end-locator; read point = marker - 1)" + placement = ( + "at the end of the last user turn (end-locator; read point = marker)" + ) elif info["type"] == ANCHOR_MODE_SR: placement = f"replacing '{sr_site}' in the generation prompt" anchor = next(iter(sr_anchors)) diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index 207c75b6..56bc138b 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -713,7 +713,6 @@ def forward( dtype=input_ids.dtype, ) is_marker = torch.isin(input_ids, marker_ids) # [batch, seq_len] - has_marker = is_marker.any(dim=1) # One verdict per request: the exit rewrites a single logit row, so a # request naming two different classifier slots has no way to report # both. Repeated markers of the SAME slot are fine -- a multi-turn @@ -731,19 +730,11 @@ def forward( ) pos = torch.arange(seq_length, device=input_ids.device) marker_idx = torch.where(is_marker, pos, -1).max(dim=1).values.clamp(min=0) - if bool((has_marker & (marker_idx == 0)).any()): - raise RuntimeError( - "Classifier marker is at position 0, so there is no last " - "content token to read (the verdict is read at marker - 1). " - "The chat template emits the marker after the last token." - ) - # A row with a marker reads its predecessor. A row without one (a plain - # LM request in a mixed batch) has no read point at all; it is parked on - # index 0 to keep the gather in range and carries slot 0, so the exit's - # is_classifier gate discards its zero verdict. - read_idx = torch.where(has_marker, marker_idx - 1, 0) + # A row without one (a plain LM request in a mixed batch) has no read + # point; the clamp parks it on index 0 and it carries slot 0, so the + # exit's is_classifier gate discards its zero verdict. batch_ar = torch.arange(batch_size, device=hidden_states.device) - req_hidden = hidden_states[batch_ar, read_idx, :] # [batch, hidden] + req_hidden = hidden_states[batch_ar, marker_idx, :] # [batch, hidden] req_slot = classifier_indices[batch_ar, marker_idx] # [batch] self._last_classifier_logits = self.classifier_head(req_hidden, req_slot) self._last_req_classifier_indices = req_slot diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index e7f4450a..dad03f0b 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -461,25 +461,21 @@ def _classifier_read_points( ) -> tuple[torch.Tensor, torch.Tensor] | None: """Per-request (hidden_state, slot_index) for the classifier verdict. - The verdict is read at the LAST CONTENT TOKEN, which the classifier - control token (marker) locates: the chat template emits the marker - immediately after the last content token, so the read point is - ``marker - 1``. + The chat template emits the marker immediately after the last content + token, so the read point is ``marker``. The marker is the LAST classifier control token in a request's slice, matched by id in ``input_ids``; ``classifier_indices`` is nonzero from any classifier marker onward, so the index alone does not say which slot a - position belongs to. The read point is ``marker - 1`` and the slot id is - read AT the marker (``marker - 1`` carries index 0). + position belongs to. Read point and slot id both come from that position. - Per request slice ``[start, end)``, three states: + Per request slice ``[start, end)``, two states: - * marker present with a predecessor in this pass -- the common case. - * marker opens the slice: its predecessor ended the previous pass, so each - pass stashes its final hidden row for a single-token look-back. - * no marker: either still mid-prefill, or the marker was in an earlier - chunk and that pass's resolved verdict is carried forward, since vLLM - consumes the verdict only on the pass it samples. + * marker present: read its row. + * no marker: the marker was in an earlier chunk, so that pass's resolved + verdict is carried forward (vLLM consumes the verdict only on the pass it + samples); or there is none at all and the request gets slot 0, whose zero + verdict the exit's ``is_classifier`` gate discards. Returns ``None`` when there is no request metadata (e.g. the startup profiling forward), so the caller emits no verdict. @@ -517,25 +513,15 @@ def _classifier_read_points( ) is_marker = torch.isin(input_ids, marker_ids) # [total_tokens] - stash = getattr(self, "_classifier_prev_pass_tail", None) or {} resolved = getattr(self, "_classifier_resolved_verdict", None) or {} req_hidden_rows = [] req_slots = [] - new_stash: dict[int, torch.Tensor] = {} new_resolved: dict[int, tuple[torch.Tensor, int]] = {} for i in range(len(qsl) - 1): start, end = qsl[i], qsl[i + 1] hits = is_marker[start:end].nonzero(as_tuple=True)[0] if hits.numel() == 0: - # No marker in this slice, which is either of two states: - # - # * the marker is still ahead (mid-prefill): keep this slice's final - # row, keyed by its global end position, in case the marker opens - # the next chunk. - # * the marker was in an earlier chunk (a rendered prompt continues - # past it with the generation prompt): carry that pass's resolved - # (row, slot) forward. vLLM consumes the verdict only on the pass - # where it samples, which is the final chunk. + # No marker in this slice carried = resolved.get(seq_lens[i] - (end - start)) if carried is not None: row, slot = carried @@ -543,7 +529,6 @@ def _classifier_read_points( req_hidden_rows.append(row) req_slots.append(slot) continue - new_stash[seq_lens[i]] = hidden_states[end - 1] req_hidden_rows.append(hidden_states[end - 1]) req_slots.append(0) continue @@ -563,29 +548,11 @@ def _classifier_read_points( marker = start + int(hits[-1]) # last id match: this slice's marker slot = int(classifier_indices[marker]) + row = hidden_states[marker] + new_resolved[seq_lens[i]] = (row, slot) req_slots.append(slot) - if marker > start: - # Common case: predecessor is in this pass. - row = hidden_states[marker - 1] - new_resolved[seq_lens[i]] = (row, slot) - req_hidden_rows.append(row) - continue - - # Edge case: marker is the first token of its slice - marker_global = seq_lens[i] - (end - start) - prev = stash.get(marker_global) - if prev is None: - raise RuntimeError( - "Classifier verdict cannot locate its read point: the marker " - f"is the first token of request {i}'s chunk (global position " - f"{marker_global}), so the last content token (marker - 1) was " - "computed in the previous pass, and no stashed hidden row is " - "available for it." - ) - new_resolved[seq_lens[i]] = (prev, slot) - req_hidden_rows.append(prev) + req_hidden_rows.append(row) - self._classifier_prev_pass_tail = new_stash self._classifier_resolved_verdict = new_resolved return ( torch.stack(req_hidden_rows, dim=0), diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py index a5eff866..d0eff6d2 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_classifier_compose.py @@ -126,9 +126,7 @@ def _classifier_prompt_ids(tokenizer, control_token_id): """`` as a batch of input ids for a detect forward. End-locator layout, matching what the chat template emits: the control token - (marker) goes AFTER the last content token, so the verdict is read at - ``marker - 1`` -- the position the head is trained on. A marker at position 0 - has no content token before it and is rejected by the backend. + (marker) goes AFTER the last content token, and is itself the read point. """ body = tokenizer.encode("Is this text safe?", add_special_tokens=False) ids = [*body, control_token_id] diff --git a/tests/hf/test_classifier.py b/tests/hf/test_classifier.py index 80fd58d4..61c853e3 100644 --- a/tests/hf/test_classifier.py +++ b/tests/hf/test_classifier.py @@ -4,8 +4,7 @@ A classifier slot is an alternative to a LoRA adapter: it shares the adapter index/control-token machinery, the switch splits classifier positions out of the LoRA stream, and the classifier head reads the final hidden state to emit -ONE verdict per request. The control token is placed after the last content -token, and the verdict is read at ``marker - 1`` (the last content token). +ONE verdict per request. The verdict is read at the control token's own position. The verdict exits as a GENERATED LABEL WORD, matching the vLLM backend: the LM head rewrites the classifier request's last logit row so the vocab is -inf @@ -71,9 +70,9 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() - # Marker (251) at the end (end-locator layout); this test checks the label-word - # rewrite of the emitted row. The read position is pinned in - # test_classifier_reads_marker_minus_one_hidden_state. + # Marker (251) at the end; this test checks the label-word rewrite of the + # emitted row. The read position is pinned in + # test_classifier_reads_marker_hidden_state. input_ids = torch.tensor([[10, 20, 30, 40, 50, 60, 70, 251]]) with torch.no_grad(): out = model(input_ids=input_ids) @@ -240,17 +239,19 @@ def test_lora_generation_on_mixed_model_decodes(): assert model.model._last_classifier_logits is None -def test_classifier_reads_marker_minus_one_hidden_state(): - """The verdict is read at ``marker - 1`` (the last content token), not at the - marker itself or the sequence's last token. +def test_classifier_reads_marker_hidden_state(): + """The verdict is read at the marker. - End-locator layout: the classifier control token is placed after the last content - token, so ``classifier_indices`` is nonzero only from the marker onward (the switch - forward-fills the index causally), and the read point is ``marker - 1``. The exact - position is pinned by capturing (a) the final post-norm hidden states via a hook on - ``model.norm`` and (b) the hidden state the classifier head actually receives, then - asserting the head's input equals the post-norm hidden at - ``marker - 1`` — and differs from the marker's and the last token's. + Reading the marker's own row is what makes the read point independent of where + the chat template puts the token. The exact position is pinned by capturing (a) + the final post-norm hidden states via a hook on ``model.norm`` and (b) the + hidden state the classifier head actually receives, then asserting the head's + input equals the post-norm hidden at the marker -- and differs from + ``marker - 1``, the position this replaced. + + The marker is placed INTERIOR (not last) so that "at the marker" is + distinguishable from "at the sequence's last row": a marker-last layout would + make the two coincide and the assertion would not pin the read point. """ config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) @@ -261,10 +262,12 @@ def test_classifier_reads_marker_minus_one_hidden_state(): model.model.classifier_head.weight[1].normal_(0.0, 0.3) model.eval() - # Marker (251) at the LAST position -> read point = marker - 1 = position 4. - input_ids = torch.tensor([[10, 20, 30, 40, 60, 251]]) + # Marker (251) INTERIOR, with content after it, so the read point is distinct + # from both marker - 1 and the sequence's last row. + input_ids = torch.tensor([[10, 20, 30, 40, 60, 251, 70, 80]]) marker_pos = 5 - read_pos = marker_pos - 1 # 4, the last content token + prev_pos = marker_pos - 1 # 4, the old read point + last_pos = input_ids.shape[1] - 1 # 7, the sequence's final row captured = {} @@ -295,18 +298,21 @@ def spy_forward(x, classifier_indices): head_x = captured["head_x"] # [batch, hidden] (gathered read-point state) assert head_x.shape == (1, config.hidden_size) - # The head was handed the post-norm hidden at marker - 1, not marker or last token. - assert torch.allclose(head_x[0], normed[read_pos], atol=1e-5), ( - "classifier head did not read the last content token (marker - 1)" + # The head was handed the post-norm hidden AT the marker. + assert torch.allclose(head_x[0], normed[marker_pos], atol=1e-5), ( + "classifier head did not read the marker's own hidden state" + ) + # Not its predecessor (the read point this replaced) ... + assert not torch.allclose(head_x[0], normed[prev_pos], atol=1e-5), ( + "classifier head read marker - 1, not the marker" ) - assert not torch.allclose(head_x[0], normed[marker_pos], atol=1e-5), ( - "classifier head read the marker itself, not marker - 1" + # ... and not the sequence's last row, which the marker is deliberately not. + assert not torch.allclose(head_x[0], normed[last_pos], atol=1e-5), ( + "classifier head read the sequence's last row, not the marker" ) - # marker_pos is also the sequence's last position, so this pins that the read - # is NOT the last token. - assert read_pos != marker_pos + assert marker_pos not in (prev_pos, last_pos) - # The slot id is read AT the marker (marker - 1 carries index 0). + # The slot id is read at the marker, which is also the read point. assert int(captured["head_slot"][0]) == 2, ( "classifier slot id must be read at the marker (adapter index 2)" ) @@ -370,6 +376,10 @@ def spy(x, classifier_indices): hook.remove() normed = captured["normed"][0] - assert torch.allclose(captured["x"][0], normed[marker_pos - 1], atol=1e-5), ( - "verdict was not read at the last marker's predecessor" + # Read at the LAST marker's own position, not the first marker's (position 0). + assert torch.allclose(captured["x"][0], normed[marker_pos], atol=1e-5), ( + "verdict was not read at the last marker" + ) + assert not torch.allclose(captured["x"][0], normed[0], atol=1e-5), ( + "verdict was read at the first marker, not the last" ) diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py index ddb62e4b..7a323c69 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -13,7 +13,7 @@ so the hooks installed here land on the model the engine actually runs. The read-index hook records, per forward pass, the global position the verdict read. -``_classifier_read_points`` reads at ``marker - 1`` (a flat within-pass index); the hook +``_classifier_read_points`` reads at the marker (a flat within-pass index); the hook maps that index through ``positions``, vLLM's own per-token coordinate, so the recorded position reflects the token actually read. """ @@ -33,9 +33,8 @@ def _install_read_index_hook(): read. Wraps ``GraniteSwitchModel._classifier_read_points`` to capture, per request, the - flat within-pass index it read (``marker - 1``, or a sentinel ``-1`` when the read - came from the cross-pass stash because the marker sat at the slice start). The - ``forward`` wrapper maps that index through ``positions`` to a global position. + flat within-pass index it read (the marker). The ``forward`` wrapper maps that + index through ``positions`` to a global position. """ from granite_switch.vllm.granite_switch_model import ( GraniteSwitchForCausalLM, @@ -47,13 +46,10 @@ def _install_read_index_hook(): orig_read = GraniteSwitchModel._classifier_read_points # Per-pass scratch for request 0 (these tests use a single request): - # flat_read_idx -- within-pass flat index the verdict read (marker-1), or None. - # stash_global -- when the read came from the previous-pass stash (marker at - # the slice start), the global position of that stashed token - # (== num_computed - 1, with num_computed derived as - # seq_lens - query_len); this pass's ``positions`` cannot show - # it, so it is carried here directly. - scratch = {"flat_read_idx": None, "stash_global": None} + # flat_read_idx -- within-pass flat index the verdict read (the marker), a + # sentinel -1 when this pass merely carried a verdict resolved + # earlier, or None. + scratch = {"flat_read_idx": None} def read_points(self, classifier_indices, hidden_states, input_ids): # Snapshot before the call: it replaces _classifier_resolved_verdict with @@ -83,9 +79,9 @@ def read_points(self, classifier_indices, hidden_states, input_ids): hits = _t.isin(input_ids[start:end], mids).nonzero(as_tuple=True)[0] if hits.numel() == 0: # No marker in this slice. If a verdict was resolved in an - # earlier pass it is carried forward, and the position it was - # read at is not addressable by this pass's ``positions``; report - # it directly. Otherwise this request is still mid-prefill. + # earlier pass it is carried forward, and that position is not + # addressable by this pass's ``positions``; report it directly. + # Otherwise this request is still mid-prefill. sql = getattr(am, "seq_lens", None) pre = ( int(sql.tolist()[0]) - (end - start) @@ -94,42 +90,26 @@ def read_points(self, classifier_indices, hidden_states, input_ids): ) if pre is not None and pre in carried_before: scratch["flat_read_idx"] = -1 - scratch["stash_global"] = scratch.get("carried_global") else: scratch["flat_read_idx"] = end - 1 else: + # The marker is its own read point, always in this pass. marker = start + int(hits[-1]) - if marker > start: - # Common case: predecessor (marker-1) is in this pass. - scratch["flat_read_idx"] = marker - 1 - scratch["resolving_flat"] = marker - 1 - else: - # Edge case: marker at slice start; the read came from the - # previous-pass stash. Its global position is - # num_computed - 1, where num_computed (tokens computed - # BEFORE this pass) is seq_lens - query_len. - sql = getattr(am, "seq_lens", None) - nc0 = ( - int(sql.tolist()[0]) - (end - start) - if sql is not None - else None - ) - scratch["flat_read_idx"] = -1 - scratch["stash_global"] = nc0 - 1 if nc0 is not None else None + scratch["flat_read_idx"] = marker + scratch["resolving_flat"] = marker except Exception: # a hook must never break a forward scratch["flat_read_idx"] = None return res def fwd(self, input_ids, positions, *args, **kwargs): scratch["flat_read_idx"] = None - scratch["stash_global"] = None out = orig_fwd(self, input_ids, positions, *args, **kwargs) try: has_cls = getattr(self.model, "classifier_head", None) is not None flat = scratch["flat_read_idx"] if has_cls and positions is not None and flat is not None: if flat == -1: - read_pos = scratch["stash_global"] # from previous pass + read_pos = scratch.get("carried_global") # resolved earlier else: read_pos = int(positions[flat].item()) if scratch.get("resolving_flat") is not None: @@ -246,7 +226,7 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): ) body = tok.encode(f"Is this text safe? {words}", add_special_tokens=False) # End-locator layout, matching the chat template: the marker follows the last - # content token, and the verdict is read at marker - 1. + # content token, and is itself the read point. prompt_ids = [*body, control_id] # Same layout, but continuing past the marker the way a rendered chat prompt # does (turn close + generation prompt). The marker is then interior, so a @@ -312,7 +292,7 @@ def run_fault_sweep(ckpt, prompt_ids, out_path): """Force the classifier verdict to read positions swept across the whole prompt and record the resulting token at each. - The correct read point is ``marker - 1`` (the last content token), located inside + The correct read point is the marker itself, located inside ``_classifier_read_points``. To sweep wrong reads, that method is replaced with a forcing version: it takes the real per-request slice but reads ``hidden_states[start + forced[0]]`` and reports slot 1 (the classifier slot) so the @@ -332,7 +312,7 @@ def run_fault_sweep(ckpt, prompt_ids, out_path): n = len(prompt_ids) - forced = [None] # None => real method (marker-1); int => forced within-slice index + forced = [None] # None => real method (the marker); int => forced slice index read_global = [None] # global position the forced read landed on orig_read = GraniteSwitchModel._classifier_read_points @@ -378,7 +358,7 @@ def one(forced_idx): passes.clear() o = llm.generate([pt], sp)[0].outputs[0] if forced_idx is None: - # Correct run: read position comes from the hook (marker - 1). + # Correct run: read position comes from the hook (the marker). last = passes[-1] if passes else None rp = last["read_global_pos"] if last else None pos = int(rp) if rp is not None else None @@ -387,7 +367,7 @@ def one(forced_idx): pos = read_global[0] return {"pos": pos, "token_id": int(o.token_ids[0]), "text": o.text} - # Correct run: reads marker - 1 (global pos n-2, marker is at n-1). + # Correct run: reads the marker (global pos n-1). correct = one(None) # Faulted sweep at a fixed stride over the prompt, excluding the correct diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_classifier_chunked_prefill.py index ab15a194..7d781029 100644 --- a/tests/vllm/test_classifier_chunked_prefill.py +++ b/tests/vllm/test_classifier_chunked_prefill.py @@ -4,24 +4,23 @@ Every other vLLM classifier path is HF-only and single-pass, yet ``enable_chunked_prefill`` is on by default in serving. Two tests here: -The prompt places the classifier control token immediately after the last -content token (marker at position ``n-1``), and the verdict is read at -``marker - 1`` (the last content token, ``n-2``) inside +The prompt places the classifier control token last (marker at position ``n-1``), +and the verdict is read AT THE MARKER inside ``GraniteSwitchModel._classifier_read_points``. 1. ``test_classifier_token_stable_across_chunk_budgets`` — the same prompt served at - ``max_num_batched_tokens`` in {unchunked, 512, 256, 128, 101, 64, 48, n_prompt-1} - emits the same verdict token, and each budget's consumed verdict is read from - ``marker - 1``. ``_classifier_read_points`` locates the marker per pass and reads - its predecessor; when a chunk boundary splits the marker from its predecessor, a - single-token cross-pass stash supplies the previous pass's last hidden row. The - ``n_prompt - 1`` budget is the arm that forces that split (marker alone in the - final chunk) — the only budget here that exercises the look-back at all. + ``max_num_batched_tokens`` in {unchunked, 256, 101, 48} emits the same verdict + token, read from the marker at every budget. Since the read point is the marker's + own row, no chunk split can separate it from the pass that computes it. 2. ``test_wrong_read_corrupts_position_and_often_verdict`` — the inverse: forcing the read to wrong positions lands at a wrong global position (asserted) and usually flips the label (reported; see that test). +3. ``test_verdict_carries_when_marker_is_not_in_the_sampling_chunk`` — the one + cross-pass path left: a rendered prompt continues past the marker, so a budget can + close the prompt after it and the sampling pass holds no marker. + Not asserted: logprob or hidden-state equality across budgets. Chunked and unchunked serving differ by bf16 reduction order in the attention/KV path (a magnitude that is hardware- and version-specific), so a numeric band would be flaky. Token stability @@ -52,29 +51,12 @@ _REPO = os.path.dirname(os.path.dirname(os.path.dirname(WORKER))) BASE_MODEL = os.environ.get("CLS_TEST_BASE", "ibm-granite/granite-4.1-3b") -# One budget per distinct final-chunk shape, which is what the read point depends on: -# ``None`` disables chunked prefill (single pass); 256 leaves a wide final chunk; 48 -# splits into many passes with a narrow one; 101 leaves exactly two tokens, so the -# marker and its read point are the last pair in the chunk. The lone-marker budget -# (final chunk of one) is derived per prompt and appended by the test below. Budgets -# differing only in pass count exercise no additional path. +# One budget per distinct final-chunk shape: ``None`` disables chunked prefill +# (single pass); 256 leaves a wide final chunk; 101 a narrow one; 48 splits into many +# passes. Budgets differing only in pass count exercise no additional path. BUDGETS = [None, 256, 101, 48] -def _lone_marker_budget(n_prompt: int) -> int: - """Budget that puts the marker ALONE in the final chunk. - - The marker is the prompt's last token, so it is the first (and only) token of - the final chunk exactly when ``n_prompt % budget == 1`` -- i.e. budget - ``n_prompt - 1``, which splits the prompt into [0, n-2] then [n-1]. That is - the one layout where the read point (``marker - 1``) lands in the *previous* - pass, exercising the cross-pass look-back in - ``GraniteSwitchModel._classifier_read_points``. Computed from ``n_prompt`` so a - tokenizer change that shifts it cannot silently stop testing this case. - """ - return n_prompt - 1 - - def _env(): env = dict(os.environ) env["VLLM_USE_V1"] = "1" @@ -168,31 +150,16 @@ def _run_fault_sweep(composed): def test_classifier_token_stable_across_chunk_budgets(composed): - """Emitted verdict token identical across chunk budgets; every budget's - consumed verdict is read from ``marker - 1`` (the last content token). - - The prompt is ``[*body, control_id]`` (marker at ``n_prompt - 1``), so the - correct read point is ``marker - 1 == n_prompt - 2``. + """Emitted verdict token identical across chunk budgets, read AT THE MARKER. - Budgets include ``n_prompt - 1``, the one budget that leaves the marker ALONE - in the final chunk (``n_prompt % budget == 1``). There the read point sits in - the *previous* pass, so this arm is what exercises the cross-pass look-back; - every other budget keeps marker and read point in the same chunk.""" + The prompt is ``[*body, control_id]``, so the marker at ``n_prompt - 1`` is also + the read point. Being the marker's own row, it is computed in whichever pass + holds the marker, for any chunk shape.""" n_prompt = len(composed["prompt_ids"]) - read_point = n_prompt - 2 # marker is at n_prompt - 1; read point = marker - 1 + read_point = n_prompt - 1 # the marker itself label_ids = set(composed["label_token_ids"]) - # The lone-marker budget (n_prompt % budget == 1) is the only layout that - # splits the marker from its read point across a chunk boundary, so it is the - # one arm that exercises the cross-pass look-back. Derived from n_prompt. - lone = _lone_marker_budget(n_prompt) - budgets = [*BUDGETS, lone] - assert n_prompt % lone == 1, ( - f"budget {lone} does not leave the marker alone in the final chunk for " - f"n_prompt={n_prompt} (n_prompt % budget == {n_prompt % lone}, want 1)" - ) - - results = {b: _run(composed, b) for b in budgets} + results = {b: _run(composed, b) for b in BUDGETS} ref = results[None] assert ref["token_id"] in label_ids, ( @@ -200,11 +167,11 @@ def test_classifier_token_stable_across_chunk_budgets(composed): f"the classifier exit did not fire" ) - # Consumed verdict reads the last content token (marker - 1), for every budget. + # Consumed verdict reads the marker's own row, for every budget. for b, r in results.items(): assert r["global_read_idx"] == read_point, ( f"budget={b}: consumed verdict read from global index " - f"{r['global_read_idx']}, expected {read_point} (marker-1) " + f"{r['global_read_idx']}, expected {read_point} (the marker) " f"(passes={r['n_passes']})" ) @@ -225,10 +192,9 @@ def test_wrong_read_corrupts_position_and_often_verdict(composed): """A wrong read position corrupts the verdict. Forces the verdict to read positions strided across the prompt and records the - emitted token at each, against a correct run reading ``marker - 1`` (the last - content token). + emitted token at each, against a correct run reading the marker itself. Asserts a wrong read always lands at a wrong global position (the correct run - reads ``n_prompt - 2``; every forced read reports a different position). + reads ``n_prompt - 1``; every forced read reports a different position). The label-flip rate is reported, not asserted: a head can map two hidden states to the same argmax label, so some wrong reads produce the correct label by chance, and @@ -236,7 +202,7 @@ def test_wrong_read_corrupts_position_and_often_verdict(composed): position-sensitive so most wrong reads do flip the label in practice. """ n_prompt = len(composed["prompt_ids"]) - read_point = n_prompt - 2 # marker at n_prompt - 1; correct read = marker - 1 + read_point = n_prompt - 1 # the marker itself label_ids = set(composed["label_token_ids"]) sweep = _run_fault_sweep(composed) @@ -244,8 +210,8 @@ def test_wrong_read_corrupts_position_and_often_verdict(composed): faulted = sweep["faulted"] assert correct["pos"] == read_point, ( - f"correct run read global pos {correct['pos']}, expected last content token " - f"{read_point} (marker - 1)" + f"correct run read global pos {correct['pos']}, expected the marker " + f"{read_point}" ) assert correct["token_id"] in label_ids, ( f"correct verdict {correct['token_id']} is not a label token {label_ids}; " @@ -290,12 +256,12 @@ def test_verdict_carries_when_marker_is_not_in_the_sampling_chunk(composed): was resolved one pass earlier and has to be carried forward. The prompt is ``[*body, marker, *tail]``, so the marker sits at ``n_tail`` tokens - from the end and its read point is ``marker - 1``. + from the end and is itself the read point. """ n_prompt = len(composed["prompt_ids_trailing"]) n_tail = composed["n_tail"] marker_pos = n_prompt - n_tail - 1 - read_point = marker_pos - 1 + read_point = marker_pos label_ids = set(composed["label_token_ids"]) assert n_tail > 0, "the trailing variant must continue past the marker" From 814e3dcc809d9220a8a44d401a28cbd1f30a5ceb Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Fri, 2 Oct 2026 03:16:27 +0000 Subject: [PATCH 04/20] feat(classifier): drop the generation prompt for classifier renders A classifier reads its verdict at the marker, so no generation prompt may follow it. vLLM's chat-completions endpoint defaults add_generation_prompt to True, so relying on callers to turn it off would silently place the assistant prompt after the marker. The composed template now forces add_generation_prompt to false when adapter_name names a classifier slot; renders with other adapters or no adapter still honor the flag. test_classifier_drops_generation_prompt pins this on granite-4.0-micro, granite-4.1-3b and granite-4.2-3b (ChatML). Signed-off-by: IRENEKO --- .../composer/tokenizer_setup.py | 6 ++++ tests/composer/test_classifier_compose.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 455fbf31..1cfaa083 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -1047,6 +1047,12 @@ def configure_chat_template( {%- set adapter_invocation_drop_chars = adapter_map[adapter_name]['invocation_drop_chars'] %} {%- endif %} {%- endif %} +{#- A classifier reads its verdict at the marker, so nothing may follow it: + drop the generation prompt even when the caller asks for one. -#} +{%- set classifier_mode = adapter_type == 'classifier' %} +{%- if classifier_mode %} +{%- set add_generation_prompt = false %} +{%- endif %} """ diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py index d0eff6d2..2ee9c5cf 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_classifier_compose.py @@ -405,3 +405,39 @@ def _first_lora_b_tensor(model): if "lora_B" in name and buf.dim() == 4: return buf.detach() return None + + +@pytest.mark.parametrize( + "base", + [ + "ibm-granite/granite-4.0-micro", + "ibm-granite/granite-4.1-3b", + "ibm-granite/granite-4.2-3b", # ChatML family + ], +) +def test_classifier_drops_generation_prompt(base): + """A classifier render ignores ``add_generation_prompt=True``. + + The verdict is read at the marker, so the assistant generation prompt must + never follow it. vLLM's chat endpoint defaults the flag to True, so the + template forces it off instead of relying on the caller. + """ + from granite_switch.composer.tokenizer_setup import ( + add_control_tokens, + configure_chat_template, + ) + + tokenizer = _tokenizer(base) + discovered = [(None, "safety", "classifier", None)] + add_control_tokens(tokenizer, discovered) + configure_chat_template(tokenizer, discovered) + messages = [{"role": "user", "content": "Is this text safe?"}] + + def render(gen, **kwargs): + return tokenizer.apply_chat_template( + messages, add_generation_prompt=gen, tokenize=False, **kwargs + ) + + assert render(True, adapter_name="safety") == render(False, adapter_name="safety") + # Base renders still honor the flag. + assert render(True) != render(False) From d48d172dbc84a59471c9c228778e2d83fc10bc2f Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Fri, 2 Oct 2026 17:44:09 +0000 Subject: [PATCH 05/20] feat(classifier): place the marker as the last token Append the classifier control token after the whole template instead of inside the last user turn. Since classifier renders already drop the generation prompt, a single append makes the marker the last token in both template families, replacing the per-format turn-close rewrites. The tutorial's raw-completions examples follow the new layout. Signed-off-by: IRENEKO --- .../composer/tokenizer_setup.py | 91 ++----------------- tests/composer/test_classifier_compose.py | 49 +++++++++- tutorials/guides/classifier_slots.md | 13 ++- 3 files changed, 63 insertions(+), 90 deletions(-) diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 1cfaa083..a31f8db8 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -145,8 +145,8 @@ def detect_template_format(chat_template: str | None) -> TemplateFormat | None: #: the end of the generation prompt (Shadow Residual, ``last_context_token``). ANCHOR_MODE_SR = "sr" -#: Placement mode where the control token is emitted *after* the last content -#: token of the last user turn (classifier slot, manifest ``kind``). +#: Placement mode where the control token is emitted as the very last token, +#: after the last turn's close (classifier slot, manifest ``kind``). ANCHOR_MODE_CLASSIFIER = "classifier" @@ -1018,7 +1018,7 @@ def configure_chat_template( for adapter_name, info in adapter_mapping.items(): # ``type`` is the placement mode: lora (sequence start) / alora (before # the invocation text) / sr (replaces the generation-prompt anchor) / - # classifier (end-locator after the last content token). + # classifier (very last token, after the last turn's close). if "invocation_text" in info: mapping_entries.append( f" '{adapter_name}': {{'token': '{info['token']}', " @@ -1183,30 +1183,6 @@ def configure_chat_template( {%- endif %} """ - # Classifier placement: the marker goes after the last user turn's content and - # before that turn's closing marker (``<|end_of_text|>`` / ``<|im_end|>``). - classifier_scan = ( - """{#- Classifier scan: find the last user message (the turn being classified). -#} -{%- if ns.adapter_token and ns.adapter_type == 'classifier' %} - {%- for _msg in """ - + fmt.loop_var_source - + """ %} - {%- if _msg.role == 'user' %} - {%- set ns.classifier_target_idx = loop.index0 %} - {%- endif %} - {%- endfor %} -{%- endif %} -""" - ) - # In-loop emit: at the target user message, emit the marker right before the - # turn-closing marker. ``{{- ns.adapter_token }}`` prints the control token; - # the surrounding content and turn-close are emitted by the base template. - classifier_emit = ( - "{%- if ns.adapter_token and ns.adapter_type == 'classifier' " - "and loop.index0 == ns.classifier_target_idx %}" - "{{- ns.adapter_token }}{%- endif %}" - ) - # Build the modified template modified_chat_template = adapter_map_def + adapter_lookup + base_chat_template @@ -1231,7 +1207,6 @@ def configure_chat_template( "\n adapter_invocation_text=adapter_invocation_text," "\n adapter_invocation_drop_chars=adapter_invocation_drop_chars," "\n alora_target_idx=-1," - "\n classifier_target_idx=-1," "\n skip_next_role_marker=false" "\n )" ) @@ -1318,7 +1293,6 @@ def configure_chat_template( modified_chat_template = ( modified_chat_template[:insertion_point] + alora_pass1 - + classifier_scan + modified_chat_template[insertion_point:] ) @@ -1350,56 +1324,11 @@ def configure_chat_template( + modified_chat_template[insertion_point:] ) - # Inject the classifier emit at the end of the last user turn, before that - # turn's closing marker. Format-specific because the two families emit the - # turn boundary differently (granite_format fuses content and the close in - # one emission; chatml emits the close as a standalone statement). - classifier_names = sorted( - n for n, i in adapter_mapping.items() if i["type"] == ANCHOR_MODE_CLASSIFIER - ) - if classifier_names and fmt.name == "chatml": - # Emit right before the user/system branch's standalone - # ``{{- '<|im_end|>\n' }}``. That close follows the low-effort if/else, - # so anchor on ``{%- endif %}\n{{- '<|im_end|>\n' }}`` scoped to the - # user branch: the FIRST ``<|im_end|>`` emission that is preceded by an - # ``{%- endif %}`` is the user/system one (assistant/tool closes are not - # preceded by an if/else over content). - chatml_user_close_re = re.compile( - r"(\{%-\s*endif\s*%\}\s*)(\{\{-\s*'<\|im_end\|>\\n'\s*\}\})" - ) - m = chatml_user_close_re.search(modified_chat_template) - if m: - insertion_point = m.start(2) - modified_chat_template = ( - modified_chat_template[:insertion_point] - + classifier_emit - + modified_chat_template[insertion_point:] - ) - elif classifier_names: - # granite_format: split the fused user emission so the marker prints - # after ``content.val`` and before ``'<|end_of_text|>\n'``. Rewrite - # ``... + content.val + '<|end_of_text|>\n' }}`` into - # ``... + content.val }}{{- '<|end_of_text|>\n' }}``. - gf_user_emit = ( - "'<|start_of_role|>' + message.role + '<|end_of_role|>' " - "+ content.val + '<|end_of_text|>\\n'" - ) - gf_user_emit_split = ( - "'<|start_of_role|>' + message.role + '<|end_of_role|>' " - "+ content.val }}" + classifier_emit + "{{- '<|end_of_text|>\\n'" - ) - if gf_user_emit in modified_chat_template: - modified_chat_template = modified_chat_template.replace( - gf_user_emit, gf_user_emit_split, 1 - ) - - if classifier_names and classifier_emit not in modified_chat_template: - raise ValueError( - f"Classifier slot(s) {classifier_names} were requested, but the " - f"classifier control-token emit could not be injected into the " - f"{fmt.name} chat template: the user-turn closing anchor was not found. " - f"Update the anchor in configure_chat_template to match this base's " - f"template." + # Classifier placement: the marker is the very last token, after the last + # turn's close. + if any(i["type"] == ANCHOR_MODE_CLASSIFIER for i in adapter_mapping.values()): + modified_chat_template += ( + "\n{%- if classifier_mode %}{{- adapter_token }}{%- endif %}" ) # Insert alora fallback before generation prompt @@ -1506,9 +1435,7 @@ def _split_merged(match: "re.Match") -> str: print(f"Chat template configured with {len(adapter_mapping)} adapter mappings:") for adapter_name, info in adapter_mapping.items(): if info["type"] == ANCHOR_MODE_CLASSIFIER: - placement = ( - "at the end of the last user turn (end-locator; read point = marker)" - ) + placement = "after the last turn's close (last token; read point = marker)" elif info["type"] == ANCHOR_MODE_SR: placement = f"replacing '{sr_site}' in the generation prompt" anchor = next(iter(sr_anchors)) diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py index 2ee9c5cf..b7482786 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_classifier_compose.py @@ -125,8 +125,8 @@ def _winning_head(num_labels, hidden, winner, dtype=torch.float32): def _classifier_prompt_ids(tokenizer, control_token_id): """`` as a batch of input ids for a detect forward. - End-locator layout, matching what the chat template emits: the control token - (marker) goes AFTER the last content token, and is itself the read point. + The control token (marker) is the last token, as in a chat-template render, + and is itself the read point. """ body = tokenizer.encode("Is this text safe?", add_special_tokens=False) ids = [*body, control_token_id] @@ -189,7 +189,7 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): # A control token id for the single classifier slot. Any spare vocab id # works for construction; the switch fires on this id, which the prompt - # helper places after the content (end-locator layout). + # helper places last, as the chat template does. control_id = tokenizer.vocab_size - 5 model = GraniteSwitchComposer.from_base_and_adapters( @@ -441,3 +441,46 @@ def render(gen, **kwargs): assert render(True, adapter_name="safety") == render(False, adapter_name="safety") # Base renders still honor the flag. assert render(True) != render(False) + + +@pytest.mark.parametrize( + "base", + [ + "ibm-granite/granite-4.0-micro", + "ibm-granite/granite-4.1-3b", + "ibm-granite/granite-4.2-3b", # ChatML family + ], +) +def test_classifier_marker_is_last_token(base): + """The marker is the very last token, after the last turn's close. + + Everything before it is byte-identical to a no-adapter render without a + generation prompt, so the decoder sees exactly the base transcript. + """ + from granite_switch.composer.tokenizer_setup import ( + add_control_tokens, + configure_chat_template, + ) + + tokenizer = _tokenizer(base) + discovered = [(None, "safety", "classifier", None)] + (control_id,), _ = add_control_tokens(tokenizer, discovered) + configure_chat_template(tokenizer, discovered) + messages = [ + {"role": "user", "content": "Is this text safe?"}, + {"role": "assistant", "content": "It looks fine."}, + {"role": "user", "content": "And this one?"}, + ] + + plain = tokenizer.apply_chat_template( + messages, add_generation_prompt=False, tokenize=False + ) + ids = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, adapter_name="safety" + ) + if not isinstance(ids, list): # BatchEncoding on newer transformers + ids = ids["input_ids"] + + assert ids[-1] == control_id + assert ids.count(control_id) == 1 + assert tokenizer.decode(ids[:-1]) == plain diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index a4077dcc..ad2581c3 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -236,15 +236,17 @@ python -m vllm.entrypoints.openai.api_server \ --port 8000 ``` -Fire a classifier by placing its control token in the prompt — one token per slot, `<|safety|>` for -the `safety` slot, `<|faithfulness|>` for the other. The completion is the emitted label word: +Fire a classifier by placing its control token at the very end of the prompt, after the last turn's +close — one token per slot, `<|safety|>` for the `safety` slot, `<|faithfulness|>` for the other. The +verdict is read at the control token, so leave out the generation prompt. The completion is the +emitted label word: ```bash curl http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -d '{ "model": "./granite-switch-detect", - "prompt": "<|safety|>Is this text safe? The weather is nice today.", + "prompt": "<|start_of_role|>user<|end_of_role|>Is this text safe? The weather is nice today.<|end_of_text|>\n<|safety|>", "max_tokens": 1 }' ``` @@ -261,7 +263,8 @@ two labels for a 2-label slot, three for a 3-label slot. ### Firing via the chat template (optional) Instead of placing `<|safety|>` by hand, you can let the composed model's chat template place it for -you by passing `adapter_name` — a classifier slot is activated exactly like any other adapter: +you by passing `adapter_name` — a classifier slot is activated exactly like any other adapter. The +template appends the control token after the last turn: ```bash curl http://localhost:8000/v1/chat/completions \ @@ -283,7 +286,7 @@ curl http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -d '{ "model": "./granite-switch-detect", - "prompt": "<|safety|>Is this text safe? ...", + "prompt": "<|start_of_role|>user<|end_of_role|>Is this text safe? ...<|end_of_text|>\n<|safety|>", "max_tokens": 1, "logprobs": 5 }' From f933c504b48001bf3fd02a7a6bc91890999ad17e Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Sat, 3 Oct 2026 03:54:55 +0000 Subject: [PATCH 06/20] refactor(classifier): compute the vLLM verdict only at marker tokens With the marker as the prompt's last token, it always lands in the final prefill chunk, so the cross-pass verdict carry is unnecessary. Run the head only on a pass's marker tokens and rewrite just their logit rows; raise when a marker is not its request's last token. Drop the read-position tests that no longer apply. Signed-off-by: IRENEKO --- .../vllm/granite_switch_model.py | 141 +++------- .../_classifier_chunked_prefill_worker.py | 264 +----------------- tests/vllm/test_classifier_chunked_prefill.py | 187 ++----------- tests/vllm/test_classifier_read_points.py | 93 ++++++ 4 files changed, 157 insertions(+), 528 deletions(-) create mode 100644 tests/vllm/test_classifier_read_points.py diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index dad03f0b..51d3c456 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -458,28 +458,23 @@ def _classifier_read_points( classifier_indices: torch.Tensor, hidden_states: torch.Tensor, input_ids: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor] | None: - """Per-request (hidden_state, slot_index) for the classifier verdict. + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None: + """``(rows, hidden, slots, num_reqs)`` for this pass's classifier markers. - The chat template emits the marker immediately after the last content - token, so the read point is ``marker``. - - The marker is the LAST classifier control token in a request's slice, - matched by id in ``input_ids``; ``classifier_indices`` is nonzero from any - classifier marker onward, so the index alone does not say which slot a - position belongs to. Read point and slot id both come from that position. - - Per request slice ``[start, end)``, two states: - - * marker present: read its row. - * no marker: the marker was in an earlier chunk, so that pass's resolved - verdict is carried forward (vLLM consumes the verdict only on the pass it - samples); or there is none at all and the request gets slot 0, whose zero - verdict the exit's ``is_classifier`` gate discards. - - Returns ``None`` when there is no request metadata (e.g. the startup - profiling forward), so the caller emits no verdict. + Each marker must be its request's last token, so it reports to that + request's logit row. ``None`` when the pass holds no marker. """ + if input_ids is None: + raise RuntimeError("Classifier requests need input_ids, not inputs_embeds.") + marker_ids = torch.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + is_marker = torch.isin(input_ids, marker_ids) # [total_tokens] + if not bool(is_marker.any()): + return None + from vllm.forward_context import get_forward_context attn_metadata = get_forward_context().attn_metadata @@ -493,71 +488,17 @@ def _classifier_read_points( if query_start_loc is None: return None - qsl = query_start_loc.tolist() - seq_lens = attn_metadata.seq_lens.tolist() - - # A request's marker is the LAST classifier control token in its slice: - # ``classifier_indices`` is nonzero from any classifier marker onward, so - # with several classifier slots the index alone does not say which slot's - # marker a position belongs to. - if input_ids is None: + pos = is_marker.nonzero(as_tuple=True)[0] # [num_markers] + ends = query_start_loc[1:].long() + # The request whose slice [start, end) holds each marker = its logit row. + rows = torch.searchsorted(ends, pos, right=True) + misplaced = pos != ends[rows] - 1 + if bool(misplaced.any()): + reqs = sorted(set(rows[misplaced].tolist())) raise RuntimeError( - "A classifier request needs input_ids to locate its control token, " - "but only inputs_embeds was provided. Classifier slots classify a " - "text prompt; pass input_ids." + f"Classifier request(s) {reqs}: control token is not their last token." ) - marker_ids = torch.tensor( - self.config.classifier_control_token_ids, - device=input_ids.device, - dtype=input_ids.dtype, - ) - is_marker = torch.isin(input_ids, marker_ids) # [total_tokens] - - resolved = getattr(self, "_classifier_resolved_verdict", None) or {} - req_hidden_rows = [] - req_slots = [] - new_resolved: dict[int, tuple[torch.Tensor, int]] = {} - for i in range(len(qsl) - 1): - start, end = qsl[i], qsl[i + 1] - hits = is_marker[start:end].nonzero(as_tuple=True)[0] - if hits.numel() == 0: - # No marker in this slice - carried = resolved.get(seq_lens[i] - (end - start)) - if carried is not None: - row, slot = carried - new_resolved[seq_lens[i]] = (row, slot) - req_hidden_rows.append(row) - req_slots.append(slot) - continue - req_hidden_rows.append(hidden_states[end - 1]) - req_slots.append(0) - continue - - # One verdict per request: the exit rewrites a single logit row, so a - # request naming two different classifier slots has no way to report - # both. Repeated markers of the SAME slot are fine (a multi-turn prompt - # carries them from earlier turns); the last one wins. - distinct = torch.unique(input_ids[start:end][hits]) - if distinct.numel() > 1: - raise RuntimeError( - "A classifier request may name only one classifier slot; " - f"request {i} carries control tokens " - f"{sorted(int(t) for t in distinct)}. The verdict exit rewrites " - "one logit row, so only one slot can report." - ) - - marker = start + int(hits[-1]) # last id match: this slice's marker - slot = int(classifier_indices[marker]) - row = hidden_states[marker] - new_resolved[seq_lens[i]] = (row, slot) - req_slots.append(slot) - req_hidden_rows.append(row) - - self._classifier_resolved_verdict = new_resolved - return ( - torch.stack(req_hidden_rows, dim=0), - torch.tensor(req_slots, dtype=torch.long, device=hidden_states.device), - ) + return rows, hidden_states[pos], classifier_indices[pos].long(), len(ends) @MULTIMODAL_REGISTRY.register_processor( @@ -760,11 +701,11 @@ def forward( classifier_indices, hidden_states, input_ids ) if read is not None: - req_hidden, req_classifier_indices = read + rows, marker_hidden, slots, num_reqs = read verdict = classifier_head( - req_hidden, req_classifier_indices - ) # [num_reqs, num_labels] - self._pending_verdict = (verdict, req_classifier_indices) + marker_hidden, slots + ) # [num_markers, max_labels] + self._pending_verdict = (rows, verdict, slots, num_reqs) return hidden_states @@ -802,30 +743,22 @@ def _apply_classifier_verdict( return logits # Consume it so a same-step prompt_logprobs pass re-applies nothing. self._pending_verdict = None - verdict, req_indices = pending - - per_slot_label_ids = self.config.classifier_label_token_ids + rows, verdict, slots, num_reqs = pending - if verdict.shape[0] != logits.shape[0]: + if logits.shape[0] != num_reqs: raise RuntimeError( - f"classifier verdict rows ({verdict.shape[0]}) do not match " - f"logits rows ({logits.shape[0]}); per-request alignment broke." + f"logits rows ({logits.shape[0]}) do not match the forward's " + f"{num_reqs} requests; per-request alignment broke." ) - is_classifier = req_indices > 0 # [num_reqs] - if not bool(is_classifier.any()): - return logits - - # Each request writes only its own slot's label ids and only that slot's - # real label count (verdict[:, :n]); padded columns are never read. - rows = is_classifier.nonzero(as_tuple=True)[0] # [num_classifier_reqs] + # Each marker writes only its own slot's label ids and only that slot's + # real label count (verdict[i, :n]); padded columns are never read. + per_slot_label_ids = self.config.classifier_label_token_ids logits[rows] = float("-inf") - for r in rows.tolist(): - slot = int(req_indices[r]) + for i, (r, slot) in enumerate(zip(rows.tolist(), slots.tolist())): ids = per_slot_label_ids[slot - 1] - n = len(ids) label_ids = torch.tensor(ids, dtype=torch.long, device=logits.device) - logits[r, label_ids] = verdict[r, :n].to(logits.dtype) + logits[r, label_ids] = verdict[i, : len(ids)].to(logits.dtype) return logits def sample( diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py index 7a323c69..f0e9ac5c 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -4,18 +4,8 @@ Each sub-command is a separate process invocation so only one vLLM engine is ever resident (same discipline as ``_generation_equivalence_worker``):: - python worker.py compose # CPU only; prints {"result": meta} - python worker.py compose-multilabel # CPU only; multi-label head + python worker.py compose-multilabel # CPU only; prints {"result": meta} python worker.py run # one engine, one budget - python worker.py fault-sweep # forced-read sweep over the prompt - -The engine core runs in-process (the driver sets ``VLLM_ENABLE_V1_MULTIPROCESSING=0``) -so the hooks installed here land on the model the engine actually runs. - -The read-index hook records, per forward pass, the global position the verdict read. -``_classifier_read_points`` reads at the marker (a flat within-pass index); the hook -maps that index through ``positions``, vLLM's own per-token coordinate, so the recorded -position reflects the token actually read. """ import json @@ -28,109 +18,6 @@ sys.path.insert(0, _REPO) -def _install_read_index_hook(): - """Record, per forward pass, the global position the classifier verdict actually - read. - - Wraps ``GraniteSwitchModel._classifier_read_points`` to capture, per request, the - flat within-pass index it read (the marker). The ``forward`` wrapper maps that - index through ``positions`` to a global position. - """ - from granite_switch.vllm.granite_switch_model import ( - GraniteSwitchForCausalLM, - GraniteSwitchModel, - ) - - passes = [] - orig_fwd = GraniteSwitchForCausalLM.forward - orig_read = GraniteSwitchModel._classifier_read_points - - # Per-pass scratch for request 0 (these tests use a single request): - # flat_read_idx -- within-pass flat index the verdict read (the marker), a - # sentinel -1 when this pass merely carried a verdict resolved - # earlier, or None. - scratch = {"flat_read_idx": None} - - def read_points(self, classifier_indices, hidden_states, input_ids): - # Snapshot before the call: it replaces _classifier_resolved_verdict with - # this pass's map, and the carry check needs the previous pass's. - carried_before = dict(getattr(self, "_classifier_resolved_verdict", None) or {}) - res = orig_read(self, classifier_indices, hidden_states, input_ids) - try: - from vllm.forward_context import get_forward_context - - am = get_forward_context().attn_metadata - if isinstance(am, list): - am = am[0] if am else None - if isinstance(am, dict): - am = next(iter(am.values()), None) - qsl = getattr(am, "query_start_loc", None) - if qsl is not None: - qsl = qsl.tolist() - start, end = qsl[0], qsl[1] # request 0 - # The marker is the last control-token id match in this slice. - import torch as _t - - mids = _t.tensor( - self.config.classifier_control_token_ids, - device=input_ids.device, - dtype=input_ids.dtype, - ) - hits = _t.isin(input_ids[start:end], mids).nonzero(as_tuple=True)[0] - if hits.numel() == 0: - # No marker in this slice. If a verdict was resolved in an - # earlier pass it is carried forward, and that position is not - # addressable by this pass's ``positions``; report it directly. - # Otherwise this request is still mid-prefill. - sql = getattr(am, "seq_lens", None) - pre = ( - int(sql.tolist()[0]) - (end - start) - if sql is not None - else None - ) - if pre is not None and pre in carried_before: - scratch["flat_read_idx"] = -1 - else: - scratch["flat_read_idx"] = end - 1 - else: - # The marker is its own read point, always in this pass. - marker = start + int(hits[-1]) - scratch["flat_read_idx"] = marker - scratch["resolving_flat"] = marker - except Exception: # a hook must never break a forward - scratch["flat_read_idx"] = None - return res - - def fwd(self, input_ids, positions, *args, **kwargs): - scratch["flat_read_idx"] = None - out = orig_fwd(self, input_ids, positions, *args, **kwargs) - try: - has_cls = getattr(self.model, "classifier_head", None) is not None - flat = scratch["flat_read_idx"] - if has_cls and positions is not None and flat is not None: - if flat == -1: - read_pos = scratch.get("carried_global") # resolved earlier - else: - read_pos = int(positions[flat].item()) - if scratch.get("resolving_flat") is not None: - # Remember where the verdict was resolved, so a later pass that - # merely carries it can report the same global position. - scratch["carried_global"] = int( - positions[scratch["resolving_flat"]].item() - ) - scratch["resolving_flat"] = None - passes.append( - {"read_global_pos": read_pos, "n_query": int(positions.shape[0])} - ) - except Exception: # a hook must never break a forward - pass - return out - - GraniteSwitchModel._classifier_read_points = read_points - GraniteSwitchForCausalLM.forward = fwd - return passes - - # Single-token label candidates, in priority order (only single-token ones are kept). _LABEL_CANDIDATES = [ "safe", @@ -152,9 +39,7 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): """Compose a single-classifier-slot checkpoint. ``num_labels`` single-token labels are taken from ``_LABEL_CANDIDATES`` and - ``scale`` is the std of the random head weight. Both are set high enough, over a - non-repetitive prompt, that the argmax varies across read positions -- which is - what makes a wrong read observable in the fault sweep. + ``scale`` is the std of the random head weight. """ import torch from safetensors.torch import save_file @@ -212,7 +97,6 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): model.save_pretrained(ckpt) tok.save_pretrained(ckpt) - # Non-repetitive body so mid-prompt hidden states differ across positions. words = ( "machine translation quietly reshaped how distant villages argued " "about rainfall while engineers debated whether copper cables or " @@ -224,35 +108,27 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): "owed to the moon each evening without fail or apparent complaint " "although sailors insisted otherwise over cheap wine and louder songs" ) - body = tok.encode(f"Is this text safe? {words}", add_special_tokens=False) - # End-locator layout, matching the chat template: the marker follows the last - # content token, and is itself the read point. - prompt_ids = [*body, control_id] - # Same layout, but continuing past the marker the way a rendered chat prompt - # does (turn close + generation prompt). The marker is then interior, so a - # chunk budget can close the prompt after it -- the pass that samples holds no - # marker at all and must consume the verdict resolved in an earlier pass. - tail = tok.encode( - "\n<|start_of_role|>assistant<|end_of_role|>", add_special_tokens=False + # The chat template's classifier layout: the user turn, its close, then the + # marker as the very last token, which is itself the read point. + body = tok.encode( + f"<|start_of_role|>user<|end_of_role|>Is this text safe? {words}" + "<|end_of_text|>\n", + add_special_tokens=False, ) - prompt_ids_trailing = [*body, control_id, *tail] + prompt_ids = [*body, control_id] return { "ckpt": ckpt, "labels": labels, "control_id": control_id, "label_token_ids": ids, "prompt_ids": prompt_ids, - "prompt_ids_trailing": prompt_ids_trailing, - "n_tail": len(tail), } def run_budget(ckpt, prompt_ids, budget, out_path): - import torch # noqa: F401 from vllm import LLM, SamplingParams from vllm.inputs import TokensPrompt - passes = _install_read_index_hook() kw = dict( model=ckpt, enforce_eager=True, @@ -267,118 +143,9 @@ def run_budget(ckpt, prompt_ids, budget, out_path): kw["max_num_batched_tokens"] = budget kw["enable_chunked_prefill"] = True llm = LLM(**kw) - sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=5) + sp = SamplingParams(max_tokens=1, temperature=0.0) o = llm.generate([TokensPrompt(prompt_token_ids=prompt_ids)], sp)[0].outputs[0] - - # The consumed verdict is the final pass's; single request => one read per pass. - last = passes[-1] if passes else None - global_idx = None - if last and last["read_global_pos"] is not None: - rp = last["read_global_pos"] - global_idx = int(rp[0]) if isinstance(rp, list) else int(rp) - - out = { - "budget": budget, - "token_id": int(o.token_ids[0]), - "text": o.text, - "n_passes": len(passes), - "global_read_idx": global_idx, - "n_prompt": len(prompt_ids), - } - json.dump(out, open(out_path, "w")) - - -def run_fault_sweep(ckpt, prompt_ids, out_path): - """Force the classifier verdict to read positions swept across the whole prompt - and record the resulting token at each. - - The correct read point is the marker itself, located inside - ``_classifier_read_points``. To sweep wrong reads, that method is replaced with a - forcing version: it takes the real per-request slice but reads - ``hidden_states[start + forced[0]]`` and reports slot 1 (the classifier slot) so the - verdict still fires. ``forced[0] is None`` runs the real method. - - The prefill runs unchunked (a single pass over all ``n`` tokens), so the within-pass - flat index equals the global position and a forced ``start + i`` reaches global - position ``i`` for any ``0 <= i <= n-1`` in one resident engine. A mutable - ``forced[0]`` cell, closed over by the monkeypatch, changes the forced index between - ``generate`` calls without rebuilding the engine.""" - import torch - from vllm import LLM, SamplingParams - from vllm.forward_context import get_forward_context - from vllm.inputs import TokensPrompt - - from granite_switch.vllm.granite_switch_model import GraniteSwitchModel - - n = len(prompt_ids) - - forced = [None] # None => real method (the marker); int => forced slice index - read_global = [None] # global position the forced read landed on - orig_read = GraniteSwitchModel._classifier_read_points - - def patched_read(self, classifier_indices, hidden_states, input_ids): - if forced[0] is None: - return orig_read(self, classifier_indices, hidden_states, input_ids) - am = get_forward_context().attn_metadata - if isinstance(am, list): - am = am[0] if am else None - if isinstance(am, dict): - am = next(iter(am.values()), None) - qsl = getattr(am, "query_start_loc", None) - if qsl is None: - return None - qsl = qsl.tolist() - start = qsl[0] # single request in these tests - flat = start + int(forced[0]) - # Unchunked single pass => flat within-pass index == global position. - read_global[0] = int(forced[0]) - req_hidden = hidden_states[flat : flat + 1] # [1, hidden] - # Slot 1 = the classifier slot, so the verdict fires like a real read. - req_slot = torch.ones(1, dtype=torch.long, device=hidden_states.device) - return req_hidden, req_slot - - GraniteSwitchModel._classifier_read_points = patched_read - try: - passes = _install_read_index_hook() # wraps the (now patched) method + forward - llm = LLM( - model=ckpt, - enforce_eager=True, - gpu_memory_utilization=0.55, - max_model_len=4096, - enable_prefix_caching=False, - dtype="bfloat16", - enable_chunked_prefill=False, - ) - sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=5) - pt = TokensPrompt(prompt_token_ids=prompt_ids) - - def one(forced_idx): - forced[0] = forced_idx - read_global[0] = None - passes.clear() - o = llm.generate([pt], sp)[0].outputs[0] - if forced_idx is None: - # Correct run: read position comes from the hook (the marker). - last = passes[-1] if passes else None - rp = last["read_global_pos"] if last else None - pos = int(rp) if rp is not None else None - else: - # Forced run: the global position is the forced index directly. - pos = read_global[0] - return {"pos": pos, "token_id": int(o.token_ids[0]), "text": o.text} - - # Correct run: reads the marker (global pos n-1). - correct = one(None) - - # Faulted sweep at a fixed stride over the prompt, excluding the correct - # read point (n-2) so every forced read is genuinely a wrong position. - step = max(1, (n - 1) // 10) - sweep_idx = sorted({i for i in range(0, n - 1, step)} - {n - 2}) - faulted = [one(i) for i in sweep_idx] - finally: - GraniteSwitchModel._classifier_read_points = orig_read - - out = {"correct": correct, "faulted": faulted, "n_prompt": n} + out = {"budget": budget, "token_id": int(o.token_ids[0]), "text": o.text} json.dump(out, open(out_path, "w")) @@ -391,15 +158,6 @@ def main(): except Exception as e: print(json.dumps({"error": f"{e}\n{traceback.format_exc()}"})) return - if mode == "fault-sweep": - out_path = sys.argv[2] - prompt_ids = json.load(open(os.environ["CLS_PIDS"])) - try: - run_fault_sweep(os.environ["CLS_CKPT"], prompt_ids, out_path) - except Exception as e: - json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) - sys.exit(1) - return if mode == "run": budget = None if sys.argv[2] == "none" else int(sys.argv[2]) out_path = sys.argv[3] diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_classifier_chunked_prefill.py index 7d781029..4590a41e 100644 --- a/tests/vllm/test_classifier_chunked_prefill.py +++ b/tests/vllm/test_classifier_chunked_prefill.py @@ -1,34 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 """Chunked-prefill coverage for the vLLM classifier verdict exit. -Every other vLLM classifier path is HF-only and single-pass, yet -``enable_chunked_prefill`` is on by default in serving. Two tests here: - -The prompt places the classifier control token last (marker at position ``n-1``), -and the verdict is read AT THE MARKER inside -``GraniteSwitchModel._classifier_read_points``. - -1. ``test_classifier_token_stable_across_chunk_budgets`` — the same prompt served at - ``max_num_batched_tokens`` in {unchunked, 256, 101, 48} emits the same verdict - token, read from the marker at every budget. Since the read point is the marker's - own row, no chunk split can separate it from the pass that computes it. - -2. ``test_wrong_read_corrupts_position_and_often_verdict`` — the inverse: forcing the - read to wrong positions lands at a wrong global position (asserted) and usually - flips the label (reported; see that test). - -3. ``test_verdict_carries_when_marker_is_not_in_the_sampling_chunk`` — the one - cross-pass path left: a rendered prompt continues past the marker, so a budget can - close the prompt after it and the sampling pass holds no marker. +``enable_chunked_prefill`` is on by default in serving. The prompt is laid out as +the chat template renders it, with the classifier control token as the very last +token, so it always lands in the final chunk. The same prompt served at +``max_num_batched_tokens`` in {unchunked, 256, 101, 48} must emit the same label +token. Not asserted: logprob or hidden-state equality across budgets. Chunked and unchunked serving differ by bf16 reduction order in the attention/KV path (a magnitude that is -hardware- and version-specific), so a numeric band would be flaky. Token stability -still catches a wrong read position or a weight/logit corruption. +hardware- and version-specific), so a numeric band would be flaky. -Each budget runs in its own subprocess so only one vLLM engine is resident, with the -engine core in-process (``VLLM_ENABLE_V1_MULTIPROCESSING=0``) so the read-index hook -lands on the running model. +Each budget runs in its own subprocess so only one vLLM engine is resident. Requires GPU + vLLM. """ @@ -60,7 +43,6 @@ def _env(): env = dict(os.environ) env["VLLM_USE_V1"] = "1" - env["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" # in-process engine core env["TOKENIZERS_PARALLELISM"] = "false" env.setdefault("PYTHONPATH", _REPO) return env @@ -89,20 +71,17 @@ def _compose(tmp_path_factory, mode, tag): result = meta["result"] result["pids_path"] = os.path.join(outdir, "pids.json") json.dump(result["prompt_ids"], open(result["pids_path"], "w")) - result["pids_trailing_path"] = os.path.join(outdir, "pids_trailing.json") - json.dump(result["prompt_ids_trailing"], open(result["pids_trailing_path"], "w")) result["outdir"] = outdir return result @pytest.fixture(scope="module") def composed(tmp_path_factory): - """Multi-label classifier over a varied prompt: the argmax is position-sensitive, - which both tests need (stable-across-budgets and the forced-read sweep).""" + """Multi-label classifier over a varied prompt.""" return _compose(tmp_path_factory, "compose-multilabel", "cls_ckpt") -def _run_worker(composed, argv, out_name, label, timeout, pids_key="pids_path"): +def _run_worker(composed, argv, out_name, label, timeout): """Run one worker invocation in its own subprocess; return the parsed result.""" out_path = os.path.join(composed["outdir"], out_name) proc = subprocess.run( @@ -110,7 +89,7 @@ def _run_worker(composed, argv, out_name, label, timeout, pids_key="pids_path"): env={ **_env(), "CLS_CKPT": composed["ckpt"], - "CLS_PIDS": composed[pids_key], + "CLS_PIDS": composed["pids_path"], }, capture_output=True, text=True, @@ -128,36 +107,20 @@ def _run_worker(composed, argv, out_name, label, timeout, pids_key="pids_path"): return res -def _run(composed, budget, pids_key="pids_path"): +def _run(composed, budget): """Serve the prompt at one chunk budget (``None`` = chunked prefill disabled).""" tag = "none" if budget is None else str(budget) - suffix = "" if pids_key == "pids_path" else "_trailing" return _run_worker( - composed, - ["run", tag], - f"run_{tag}{suffix}.json", - f"budget={budget}{suffix}", - 1200, - pids_key=pids_key, - ) - - -def _run_fault_sweep(composed): - """Run the full-prompt forced-read sweep in one subprocess.""" - return _run_worker( - composed, ["fault-sweep"], "fault_sweep.json", "fault-sweep", 1800 + composed, ["run", tag], f"run_{tag}.json", f"budget={budget}", 1200 ) def test_classifier_token_stable_across_chunk_budgets(composed): - """Emitted verdict token identical across chunk budgets, read AT THE MARKER. - - The prompt is ``[*body, control_id]``, so the marker at ``n_prompt - 1`` is also - the read point. Being the marker's own row, it is computed in whichever pass - holds the marker, for any chunk shape.""" - n_prompt = len(composed["prompt_ids"]) - read_point = n_prompt - 1 # the marker itself + """Emitted verdict token identical across chunk budgets.""" label_ids = set(composed["label_token_ids"]) + assert len(composed["prompt_ids"]) > min(b for b in BUDGETS if b), ( + "the prompt must be longer than the smallest budget so it chunks" + ) results = {b: _run(composed, b) for b in BUDGETS} @@ -166,126 +129,8 @@ def test_classifier_token_stable_across_chunk_budgets(composed): f"unchunked verdict {ref['token_id']} is not a label token {label_ids}; " f"the classifier exit did not fire" ) - - # Consumed verdict reads the marker's own row, for every budget. - for b, r in results.items(): - assert r["global_read_idx"] == read_point, ( - f"budget={b}: consumed verdict read from global index " - f"{r['global_read_idx']}, expected {read_point} (the marker) " - f"(passes={r['n_passes']})" - ) - - # Verdict token stable across budgets. for b, r in results.items(): assert r["token_id"] == ref["token_id"], ( f"budget={b} emitted {r['token_id']} ({r['text']!r}) != unchunked " f"{ref['token_id']} ({ref['text']!r})" ) - - # The smallest budget actually chunked. - assert results[48]["n_passes"] > 1, ( - f"budget=48 ran in {results[48]['n_passes']} pass(es); expected chunking." - ) - - -def test_wrong_read_corrupts_position_and_often_verdict(composed): - """A wrong read position corrupts the verdict. - - Forces the verdict to read positions strided across the prompt and records the - emitted token at each, against a correct run reading the marker itself. - Asserts a wrong read always lands at a wrong global position (the correct run - reads ``n_prompt - 1``; every forced read reports a different position). - - The label-flip rate is reported, not asserted: a head can map two hidden states to - the same argmax label, so some wrong reads produce the correct label by chance, and - a mismatch threshold would be flaky. The multi-label head keeps the argmax - position-sensitive so most wrong reads do flip the label in practice. - """ - n_prompt = len(composed["prompt_ids"]) - read_point = n_prompt - 1 # the marker itself - label_ids = set(composed["label_token_ids"]) - - sweep = _run_fault_sweep(composed) - correct = sweep["correct"] - faulted = sweep["faulted"] - - assert correct["pos"] == read_point, ( - f"correct run read global pos {correct['pos']}, expected the marker " - f"{read_point}" - ) - assert correct["token_id"] in label_ids, ( - f"correct verdict {correct['token_id']} is not a label token {label_ids}; " - f"the classifier exit did not fire" - ) - # Expect a spread of distinct positions, every one inside 0..n-2 and none at - # the correct read point. - assert len(faulted) >= 3, ( - f"fault sweep produced only {len(faulted)} forced reads; expected a spread" - ) - positions_set = {f["pos"] for f in faulted} - assert all(0 <= p <= n_prompt - 2 for p in positions_set), ( - f"forced reads landed outside the prompt pool 0..{n_prompt - 2}: " - f"{sorted(positions_set)}" - ) - for f in faulted: - assert f["pos"] != read_point, ( - f"forced read reported global pos {f['pos']} == correct read point " - f"{read_point}; the forced index did not move the read" - ) - - # Report the label-flip rate (not asserted; see docstring). - n_mismatch = sum(1 for f in faulted if f["token_id"] != correct["token_id"]) - hist = {} - for f in faulted: - hist[f["text"]] = hist.get(f["text"], 0) + 1 - positions = sorted(f["pos"] for f in faulted) - frac = n_mismatch / len(faulted) if faulted else 0.0 - print( - f"\nfault-sweep: {n_mismatch}/{len(faulted)} wrong-read positions " - f"(global {positions[0]}..{positions[-1]}) produced a label != correct " - f"{correct['text']!r} ({frac:.0%}); labels={hist}" - ) - - -def test_verdict_carries_when_marker_is_not_in_the_sampling_chunk(composed): - """A marker in an earlier chunk still produces its verdict. - - A rendered chat prompt continues past the classifier marker (turn close plus - generation prompt), so the marker is interior. With a budget that closes a chunk - at the marker, the pass vLLM samples on contains no marker at all -- the verdict - was resolved one pass earlier and has to be carried forward. - - The prompt is ``[*body, marker, *tail]``, so the marker sits at ``n_tail`` tokens - from the end and is itself the read point. - """ - n_prompt = len(composed["prompt_ids_trailing"]) - n_tail = composed["n_tail"] - marker_pos = n_prompt - n_tail - 1 - read_point = marker_pos - label_ids = set(composed["label_token_ids"]) - assert n_tail > 0, "the trailing variant must continue past the marker" - - # Budget that ends a chunk exactly at the marker: the next pass starts at the - # first tail token and holds no marker. - budget = marker_pos + 1 - res = _run(composed, budget, pids_key="pids_trailing_path") - - assert res["n_passes"] > 1, ( - f"budget={budget} ran in {res['n_passes']} pass(es); expected the prompt to " - f"split after the marker" - ) - assert res["token_id"] in label_ids, ( - f"verdict {res['token_id']} is not a label token {label_ids}; the verdict " - f"resolved before the sampling chunk was not carried forward" - ) - assert res["global_read_idx"] == read_point, ( - f"carried verdict was read from global index {res['global_read_idx']}, " - f"expected {read_point} (marker at {marker_pos}, {n_tail} tail tokens)" - ) - - # Same prompt unchunked must agree. - ref = _run(composed, None, pids_key="pids_trailing_path") - assert res["token_id"] == ref["token_id"], ( - f"chunked verdict {res['token_id']} != unchunked {ref['token_id']} on the " - f"trailing-token prompt" - ) diff --git a/tests/vllm/test_classifier_read_points.py b/tests/vllm/test_classifier_read_points.py new file mode 100644 index 00000000..87dac3e7 --- /dev/null +++ b/tests/vllm/test_classifier_read_points.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for ``GraniteSwitchModel._classifier_read_points`` (vLLM). + +The read runs outside the compiled region and only needs ``input_ids``, the +final hidden states and ``query_start_loc``, so it is exercised here with a fake +forward context instead of a served model. Serving coverage, including chunked +prefill, lives in ``test_classifier_chunked_prefill.py``. +""" + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("vllm") + +import vllm.forward_context + +from granite_switch.vllm.granite_switch_model import GraniteSwitchModel + +# Two classifier slots: control id 90 fires slot 2, 91 fires slot 3. +_SLOT_OF = {90: 2, 91: 3} +_MODEL = SimpleNamespace( + config=SimpleNamespace(classifier_control_token_ids=list(_SLOT_OF)) +) + + +def _read(monkeypatch, slices, query_start_loc=None): + """Run the read over one pass made of the given per-request token slices.""" + ids = torch.tensor([t for s in slices for t in s]) + # The switch's index stream: a marker's slot from the marker onward, per request. + indices = [] + for s in slices: + current = 0 + for t in s: + current = _SLOT_OF.get(t, current) + indices.append(current) + classifier_indices = torch.tensor(indices) + hidden = torch.arange(len(ids), dtype=torch.float32).unsqueeze(1) # row i == i + if query_start_loc is None: + bounds = [0] + for s in slices: + bounds.append(bounds[-1] + len(s)) + query_start_loc = bounds + ctx = SimpleNamespace( + attn_metadata=SimpleNamespace(query_start_loc=torch.tensor(query_start_loc)) + ) + monkeypatch.setattr(vllm.forward_context, "get_forward_context", lambda: ctx) + return GraniteSwitchModel._classifier_read_points( + _MODEL, classifier_indices, hidden, ids + ) + + +def test_reports_each_marker_to_its_requests_row(monkeypatch): + rows, hidden, slots, num_reqs = _read(monkeypatch, [[5, 6, 90], [7, 8], [9, 91]]) + # Markers sit at flat 2 and 6, the last rows of requests 0 and 2. + assert rows.tolist() == [0, 2] + assert hidden.squeeze(1).tolist() == [2.0, 6.0] + assert slots.tolist() == [2, 3] + assert num_reqs == 3 + + +def test_pass_without_a_marker_returns_none(monkeypatch): + assert _read(monkeypatch, [[5, 6], [7]]) is None + + +def test_marker_before_the_last_token_raises(monkeypatch): + with pytest.raises(RuntimeError, match=r"request\(s\) \[1\]"): + _read(monkeypatch, [[5, 90], [90, 6, 7]]) + + +def test_two_markers_raise_even_when_one_is_last(monkeypatch): + with pytest.raises(RuntimeError, match="not their last token"): + _read(monkeypatch, [[90, 5, 91]]) + + +def test_empty_slice_does_not_shift_rows(monkeypatch): + # A padded query_start_loc repeats its final offset: the third slice is empty. + rows, _, slots, num_reqs = _read( + monkeypatch, [[5, 90], [6]], query_start_loc=[0, 2, 3, 3] + ) + assert rows.tolist() == [0] + assert slots.tolist() == [2] + assert num_reqs == 3 + + +def test_no_metadata_returns_none(monkeypatch): + ctx = SimpleNamespace(attn_metadata=None) + monkeypatch.setattr(vllm.forward_context, "get_forward_context", lambda: ctx) + out = GraniteSwitchModel._classifier_read_points( + _MODEL, torch.tensor([2]), torch.zeros(1, 1), torch.tensor([90]) + ) + assert out is None From cba1848bc2ad8815b4eee39bc43e66300baf8d50 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Sat, 3 Oct 2026 03:54:58 +0000 Subject: [PATCH 07/20] refactor(classifier): align the HF verdict with vLLM Compute the verdict only at classifier markers, require each marker to be its row's last real token, and write it to the marker's own logit row. Decode steps after the verdict now run as plain LM steps instead of raising, matching vLLM. Signed-off-by: IRENEKO --- .../hf/modeling_granite_switch.py | 144 +++++------ tests/hf/test_classifier.py | 223 +++++------------- 2 files changed, 119 insertions(+), 248 deletions(-) diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index 56bc138b..3bbd013e 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -678,69 +678,50 @@ def forward( all_hidden_states += (hidden_states,) # Classifier heads read the final post-norm hidden state (the same - # representation the LM head consumes) and emit one verdict per request - # when routed to a classifier. A detect request is prefill-only: it - # classifies the prompt and ships the verdict as the replacements to - # rewrite that request's logit row into its label word. Reading the - # final state is the initial supported read point; earlier layers can - # be added later. - - # Per-request gate so a mixed model still serves LoRA/base generation. - active_classifier = self.classifier_head is not None and bool( - (classifier_indices > 0).any() - ) - if active_classifier: - if cache_position[0] != 0: - raise RuntimeError( - "Classifier requests are detect-only (prefill, no " - "generation); got a cached decode step " - f"(cache_position starts at {int(cache_position[0])})." - ) - # The marker is the LAST classifier control token in the original - # input_ids: ``classifier_indices`` is nonzero from any classifier - # marker onward, so with several classifier slots the index alone does - # not say which slot's marker a position belongs to. A multi-turn prompt - # can also carry markers from earlier turns. + # representation the LM head consumes), only at classifier markers. Each + # marker must be its row's last real token, where the exit rewrites that + # row's logits into a label word. + self._last_classifier_verdict = None + if self.classifier_head is not None: if input_ids is None: - raise RuntimeError( - "A classifier request needs input_ids to locate its control " - "token, but only inputs_embeds was provided. Classifier slots " - "classify a text prompt; pass input_ids." - ) - marker_ids = torch.tensor( - self.config.classifier_control_token_ids, - device=input_ids.device, - dtype=input_ids.dtype, - ) - is_marker = torch.isin(input_ids, marker_ids) # [batch, seq_len] - # One verdict per request: the exit rewrites a single logit row, so a - # request naming two different classifier slots has no way to report - # both. Repeated markers of the SAME slot are fine -- a multi-turn - # prompt carries them from earlier turns and the last one wins. - marked = torch.where(is_marker, input_ids, -1) # [batch, seq_len] - for row in range(batch_size): - distinct = torch.unique(marked[row]) - distinct = distinct[distinct >= 0] - if distinct.numel() > 1: + if bool((classifier_indices > 0).any()): raise RuntimeError( - "A classifier request may name only one classifier slot; " - f"request {row} carries control tokens " - f"{sorted(int(t) for t in distinct)}. The verdict exit " - "rewrites one logit row, so only one slot can report." + "Classifier requests need input_ids, not inputs_embeds." + ) + else: + marker_ids = torch.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + is_marker = torch.isin(input_ids, marker_ids) # [batch, seq_len] + if bool(is_marker.any()): + rows, pos = is_marker.nonzero(as_tuple=True) + # Last real token of each row among this step's tokens. + if attention_mask is not None and attention_mask.dim() == 2: + real = attention_mask[:, -seq_length:].long() + cols = torch.arange(seq_length, device=real.device) + last = (real * cols).max(dim=1).values + else: + last = torch.full( + (batch_size,), seq_length - 1, device=input_ids.device + ) + misplaced = pos != last[rows] + if bool(misplaced.any()): + bad = sorted(set(rows[misplaced].tolist())) + raise RuntimeError( + f"Classifier request(s) {bad}: control token is not " + "their last token." + ) + slots = classifier_indices[rows, pos] + verdict = self.classifier_head(hidden_states[rows, pos], slots) + self._last_classifier_verdict = ( + rows, + pos, + verdict, + slots, + seq_length, ) - pos = torch.arange(seq_length, device=input_ids.device) - marker_idx = torch.where(is_marker, pos, -1).max(dim=1).values.clamp(min=0) - # A row without one (a plain LM request in a mixed batch) has no read - # point; the clamp parks it on index 0 and it carries slot 0, so the - # exit's is_classifier gate discards its zero verdict. - batch_ar = torch.arange(batch_size, device=hidden_states.device) - req_hidden = hidden_states[batch_ar, marker_idx, :] # [batch, hidden] - req_slot = classifier_indices[batch_ar, marker_idx] # [batch] - self._last_classifier_logits = self.classifier_head(req_hidden, req_slot) - self._last_req_classifier_indices = req_slot - else: - self._last_classifier_logits = None - self._last_req_classifier_indices = None if not return_dict: return tuple( @@ -882,38 +863,31 @@ def forward( ) def _apply_classifier_verdict(self, logits: torch.Tensor) -> torch.Tensor: - """Rewrite classifier requests' last logit row to emit their label word. + """Rewrite each classifier marker's logit row to emit its label word. - For each classifier request set the whole vocab to ``-inf`` at the last - prompt position, place the per-label verdict scores on each label word's - token id (``classifier_label_token_ids``) so decoding emits that word. - Non-classifier rows keep their normal LM logits. + The row is blanked to ``-inf`` and the slot's per-label verdict scores are + placed on its label token ids (``classifier_label_token_ids``). """ - verdict = getattr(self.model, "_last_classifier_logits", None) - req_slot = getattr(self.model, "_last_req_classifier_indices", None) - if verdict is None or req_slot is None: + pending = getattr(self.model, "_last_classifier_verdict", None) + if pending is None: return logits + rows, pos, verdict, slots, seq_length = pending - per_slot_label_ids = self.config.classifier_label_token_ids - - if verdict.shape[0] != logits.shape[0]: + # logits keeps the last logits.shape[1] positions (logits_to_keep). + cols = pos - (seq_length - logits.shape[1]) + if bool((cols < 0).any()): raise RuntimeError( - f"classifier verdict rows ({verdict.shape[0]}) do not match " - f"logits rows ({logits.shape[0]}); per-request alignment broke." + "A classifier marker's logits were not kept (logits_to_keep)." ) - is_classifier = req_slot > 0 # [batch] - if not bool(is_classifier.any()): - return logits - - # Each request writes only its own slot's label ids and only that slot's - # real label count (verdict[:, :n]); padded columns are never read. - rows = is_classifier.nonzero(as_tuple=True)[0] # [num_classifier_reqs] - logits[rows, -1, :] = float("-inf") - for r in rows.tolist(): - slot = int(req_slot[r]) + # Each marker writes only its own slot's label ids and only that slot's + # real label count (verdict[i, :n]); padded columns are never read. + per_slot_label_ids = self.config.classifier_label_token_ids + logits[rows, cols, :] = float("-inf") + for i, (r, c, slot) in enumerate( + zip(rows.tolist(), cols.tolist(), slots.tolist()) + ): ids = per_slot_label_ids[slot - 1] - n = len(ids) label_ids = torch.tensor(ids, dtype=torch.long, device=logits.device) - logits[r, -1, label_ids] = verdict[r, :n].to(logits.dtype) + logits[r, c, label_ids] = verdict[i, : len(ids)].to(logits.dtype) return logits diff --git a/tests/hf/test_classifier.py b/tests/hf/test_classifier.py index 61c853e3..ce7a9b30 100644 --- a/tests/hf/test_classifier.py +++ b/tests/hf/test_classifier.py @@ -4,12 +4,12 @@ A classifier slot is an alternative to a LoRA adapter: it shares the adapter index/control-token machinery, the switch splits classifier positions out of the LoRA stream, and the classifier head reads the final hidden state to emit -ONE verdict per request. The verdict is read at the control token's own position. +ONE verdict per request. The verdict is read at the control token, which must be +the request's last token. The verdict exits as a GENERATED LABEL WORD, matching the vLLM backend: the LM -head rewrites the classifier request's last logit row so the vocab is -inf -except at ``classifier_label_token_ids``, where the per-label verdict scores are -placed. There is no separate float ``classifier_logits`` output field. +head rewrites the marker's logit row so the vocab is -inf except at +``classifier_label_token_ids``, where the per-label verdict scores are placed. There is no separate float ``classifier_logits`` output field. CPU, random weights, no checkpoint — mirrors tests/hf/test_model_forward.py. """ @@ -70,9 +70,7 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() - # Marker (251) at the end; this test checks the label-word rewrite of the - # emitted row. The read position is pinned in - # test_classifier_reads_marker_hidden_state. + # Marker (251) last; this test checks the label-word rewrite of the emitted row. input_ids = torch.tensor([[10, 20, 30, 40, 50, 60, 70, 251]]) with torch.no_grad(): out = model(input_ids=input_ids) @@ -89,7 +87,7 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): assert torch.isfinite(last_row[torch.tensor(label_ids)]).all() # The label-id scores equal the classifier head's per-label verdict, sliced # to this slot's real label count (padded columns are never scattered). - verdict = model.model._last_classifier_logits[0, :n] # [n_labels] + verdict = model.model._last_classifier_verdict[2][0, :n] # [n_labels] assert torch.allclose(last_row[torch.tensor(label_ids)], verdict.to(last_row.dtype)) # The argmax (emitted token) is one of the label ids. assert int(last_row.argmax()) in label_ids @@ -122,7 +120,7 @@ def test_mixed_batch_only_classifier_row_rewritten(): input_ids = torch.tensor( [ - [10, 20, 251, 40, 50], # classifier + [10, 20, 30, 40, 251], # classifier [10, 20, 30, 40, 50], # plain ] ) @@ -139,27 +137,27 @@ def test_mixed_batch_only_classifier_row_rewritten(): assert torch.isfinite(out.logits[1, -1, :]).all() -def test_classifier_is_detect_only_decode_step_raises(): - # A classifier request is prefill-only: it classifies the prompt and stops. - # A cached decode step (generation) on a classifier model is misuse and - # must fail loud rather than produce a verdict over generated tokens. +def test_decode_step_after_verdict_is_plain_lm(): + # After the verdict token, a cached decode step carries no marker, so it runs + # as a plain LM step: no verdict, full finite vocab row. config = _classifier_config(["lora", "classifier"]) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() - input_ids = torch.tensor([[10, 20, 251, 40, 50]]) + input_ids = torch.tensor([[10, 20, 30, 251]]) with torch.no_grad(): out = model(input_ids=input_ids, use_cache=True) - assert out.logits.shape == (1, 5, config.vocab_size) + assert model.model._last_classifier_verdict is not None - with pytest.raises(RuntimeError, match="detect-only"): - with torch.no_grad(): - model( - input_ids=torch.tensor([[60]]), - past_key_values=out.past_key_values, - use_cache=True, - ) + with torch.no_grad(): + step = model( + input_ids=torch.tensor([[100]]), + past_key_values=out.past_key_values, + use_cache=True, + ) + assert model.model._last_classifier_verdict is None + assert torch.isfinite(step.logits).all() def test_classifier_missing_label_token_ids_fails_loud(): @@ -199,7 +197,7 @@ def test_lora_only_path_unaffected(): out = model(input_ids=input_ids) assert model.model.classifier_head is None - assert model.model._last_classifier_logits is None + assert model.model._last_classifier_verdict is None assert torch.isfinite(out.logits).all() assert out.logits.shape == (1, 6, config.vocab_size) @@ -207,10 +205,7 @@ def test_lora_only_path_unaffected(): def test_lora_generation_on_mixed_model_decodes(): # A model composed WITH a classifier slot also serves plain LoRA/base # generation. Firing the LoRA (adapter 1, token 250) — not the classifier — - # must be allowed to decode multiple tokens: the detect-only guard keys on an - # actually-fired classifier, not on the mere presence of a classifier head - # bank. (Regression: the guard previously raised on ANY cached decode step of - # a classifier-composed model, crashing all multi-token LoRA generation.) + # must decode multiple tokens with no verdict and full LM logits. config = _classifier_config(["lora", "classifier"]) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) @@ -222,12 +217,10 @@ def test_lora_generation_on_mixed_model_decodes(): with torch.no_grad(): out = model(input_ids=input_ids, use_cache=True) # No classifier fired -> no verdict, full finite vocab row. - assert model.model._last_classifier_logits is None + assert model.model._last_classifier_verdict is None assert torch.isfinite(out.logits).all() - # Cached decode step of the same (LoRA) request must NOT raise and must - # yield finite logits — classifier_indices stays zero, so the detect-only - # guard does not fire. + # Cached decode step of the same (LoRA) request yields finite logits. with torch.no_grad(): step = model( input_ids=torch.tensor([[60]]), @@ -236,150 +229,54 @@ def test_lora_generation_on_mixed_model_decodes(): ) assert step.logits.shape == (1, 1, config.vocab_size) assert torch.isfinite(step.logits).all() - assert model.model._last_classifier_logits is None - - -def test_classifier_reads_marker_hidden_state(): - """The verdict is read at the marker. - - Reading the marker's own row is what makes the read point independent of where - the chat template puts the token. The exact position is pinned by capturing (a) - the final post-norm hidden states via a hook on ``model.norm`` and (b) the - hidden state the classifier head actually receives, then asserting the head's - input equals the post-norm hidden at the marker -- and differs from - ``marker - 1``, the position this replaced. - - The marker is placed INTERIOR (not last) so that "at the marker" is - distinguishable from "at the sequence's last row": a marker-last layout would - make the two coincide and the assertion would not pin the read point. - """ - config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + assert model.model._last_classifier_verdict is None + + +@pytest.mark.parametrize( + "kinds, input_ids", + [ + # Content after the marker. + (["lora", "classifier"], [[10, 20, 251, 30, 40]]), + # Two different classifier slots, the second one last. + (["classifier", "classifier"], [[250, 10, 20, 30, 251]]), + # The same slot's marker twice. + (["lora", "classifier"], [[251, 10, 20, 30, 251]]), + ], + ids=["content-after", "two-slots", "repeated"], +) +def test_marker_that_is_not_the_last_token_raises(kinds, input_ids): + """The verdict is read at the last token, so any other marker would be + ignored silently; it raises instead.""" + config = _classifier_config(kinds, label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) - # Populate the classifier bank (slot 1 = adapter index 2) so the head is a real, - # non-zero map (a zero bank would make every position's verdict identical). - with torch.no_grad(): - model.model.classifier_head.weight[1].normal_(0.0, 0.3) model.eval() - # Marker (251) INTERIOR, with content after it, so the read point is distinct - # from both marker - 1 and the sequence's last row. - input_ids = torch.tensor([[10, 20, 30, 40, 60, 251, 70, 80]]) - marker_pos = 5 - prev_pos = marker_pos - 1 # 4, the old read point - last_pos = input_ids.shape[1] - 1 # 7, the sequence's final row - - captured = {} - - # Capture the final post-norm hidden states (exactly what the classifier reads). - def norm_hook(_module, _inp, output): - captured["normed"] = output.detach().clone() - - h = model.model.norm.register_forward_hook(norm_hook) - - # Spy on the head input so we see which hidden state it was handed. - real_head = model.model.classifier_head - orig_forward = real_head.forward - - def spy_forward(x, classifier_indices): - captured["head_x"] = x.detach().clone() - captured["head_slot"] = classifier_indices.detach().clone() - return orig_forward(x, classifier_indices) - - real_head.forward = spy_forward - try: + with pytest.raises(RuntimeError, match="not their last token"): with torch.no_grad(): - model(input_ids=input_ids) - finally: - real_head.forward = orig_forward - h.remove() - - normed = captured["normed"][0] # [seq_len, hidden] - head_x = captured["head_x"] # [batch, hidden] (gathered read-point state) - assert head_x.shape == (1, config.hidden_size) - - # The head was handed the post-norm hidden AT the marker. - assert torch.allclose(head_x[0], normed[marker_pos], atol=1e-5), ( - "classifier head did not read the marker's own hidden state" - ) - # Not its predecessor (the read point this replaced) ... - assert not torch.allclose(head_x[0], normed[prev_pos], atol=1e-5), ( - "classifier head read marker - 1, not the marker" - ) - # ... and not the sequence's last row, which the marker is deliberately not. - assert not torch.allclose(head_x[0], normed[last_pos], atol=1e-5), ( - "classifier head read the sequence's last row, not the marker" - ) - assert marker_pos not in (prev_pos, last_pos) - - # The slot id is read at the marker, which is also the read point. - assert int(captured["head_slot"][0]) == 2, ( - "classifier slot id must be read at the marker (adapter index 2)" - ) - - -def test_two_classifier_slots_in_one_request_rejected(): - """A request may name only one classifier slot. - - The verdict exit rewrites a single logit row, so two slots cannot both report. - ``classifier_indices`` is nonzero from either marker onward, so silently - taking one would compute the verdict from whichever position won the race. - """ - config = _classifier_config( - ["classifier", "classifier"], label_token_ids=(100, 200) - ) - model = GraniteSwitchForCausalLM(config) - _set_adapter_token_ids(model, config.adapter_token_ids) - model.eval() + model(input_ids=torch.tensor(input_ids)) - first_tok, last_tok = config.adapter_token_ids - with pytest.raises(RuntimeError, match="only one classifier slot"): - with torch.no_grad(): - model(input_ids=torch.tensor([[first_tok, 10, 20, 30, last_tok]])) - - -def test_repeated_marker_for_one_slot_reads_the_last(): - """The same slot's marker may repeat; the last occurrence is the read point. - A multi-turn prompt carries markers from earlier turns, so a repeat is - legitimate and must not be rejected -- only the current turn's marker decides - where the verdict is read. - """ +def test_right_padded_row_rewrites_the_marker_column(): + """With right padding the marker is the row's last real token, not its last + column; the verdict lands on the marker's own logit row.""" config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) - with torch.no_grad(): - model.model.classifier_head.weight[1].normal_(0.0, 0.3) model.eval() - cls_tok = config.adapter_token_ids[1] - input_ids = torch.tensor([[cls_tok, 10, 20, 30, cls_tok]]) - marker_pos = 4 - - captured = {} - real_head = model.model.classifier_head - orig_forward = real_head.forward + input_ids = torch.tensor([[10, 20, 251, 0, 0], [10, 20, 30, 40, 251]]) + attention_mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]) + with torch.no_grad(): + out = model(input_ids=input_ids, attention_mask=attention_mask) - def spy(x, classifier_indices): - captured["x"] = x.detach().clone() - return orig_forward(x, classifier_indices) + label_ids = config.classifier_label_token_ids[1] + non_label = [t for t in range(config.vocab_size) if t not in label_ids] + assert torch.isinf(out.logits[0, 2, non_label]).all() + assert torch.isfinite(out.logits[0, -1, :]).all() # padding untouched + assert torch.isinf(out.logits[1, -1, non_label]).all() - real_head.forward = spy - hook = model.model.norm.register_forward_hook( - lambda _m, _i, out: captured.__setitem__("normed", out.detach().clone()) - ) - try: + # Keeping only the last position drops row 0's marker row: fail loud. + with pytest.raises(RuntimeError, match="logits_to_keep"): with torch.no_grad(): - model(input_ids=input_ids) - finally: - real_head.forward = orig_forward - hook.remove() - - normed = captured["normed"][0] - # Read at the LAST marker's own position, not the first marker's (position 0). - assert torch.allclose(captured["x"][0], normed[marker_pos], atol=1e-5), ( - "verdict was not read at the last marker" - ) - assert not torch.allclose(captured["x"][0], normed[0], atol=1e-5), ( - "verdict was read at the first marker, not the last" - ) + model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1) From 7d05aa43517ed1b9ce940f59313e19e48d764309 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Sat, 3 Oct 2026 05:28:20 +0000 Subject: [PATCH 08/20] feat(classifier): exchange the classifier token for the role-open marker Give classifier slots their own substitute instead of inheriting the LoRA probe: the base's role-open marker, which follows the last turn's close where the classifier token sits. The tutorial tells trainers to read the hidden state at that marker. Signed-off-by: IRENEKO --- .../composer/compose_granite_switch.py | 17 ++++++++++++++ .../composer/tokenizer_setup.py | 7 ++++++ tests/composer/test_base_reset_token.py | 21 +++++++++++++++--- tests/composer/test_classifier_compose.py | 22 +++++++++++++++++++ tests/composer/test_tokenizer_setup.py | 15 +++++++++++++ tutorials/guides/classifier_slots.md | 7 +++++- 6 files changed, 85 insertions(+), 4 deletions(-) diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index 58aec734..d620b93b 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -76,6 +76,7 @@ build_substitute_token_ids, configure_audio_chat_template, configure_chat_template, + detect_template_format, find_reserved_never_emitted_token_id, load_activation_anchor, resolve_label_token_ids, @@ -121,10 +122,26 @@ def build_control_token_lists(tokenizer, all_discovered, base_reset): all_discovered, _probe_lora_substitute_token_id(tokenizer), base_reset=base_reset, + classifier_substitute_id=_classifier_substitute_token_id(tokenizer), ) return adapter_token_ids, special_tokens, adapter_substitute_token_ids +def _classifier_substitute_token_id(tokenizer) -> int: + """The base's role-open marker id, which follows the last turn's close.""" + fmt = detect_template_format(tokenizer.chat_template) + if fmt is None: + raise ValueError( + "Unrecognized chat template; cannot pick the classifier substitute." + ) + ids = tokenizer.encode(fmt.role_open_marker, add_special_tokens=False) + if len(ids) != 1: + raise ValueError( + f"Role-open marker {fmt.role_open_marker!r} encodes to {len(ids)} tokens." + ) + return ids[0] + + def _probe_lora_substitute_token_id(tokenizer) -> int: """Ask the tokenizer which token naturally appears at sequence position 0 of a rendered no-adapter chat. diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index a31f8db8..7dc83dee 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -836,6 +836,7 @@ def build_substitute_token_ids( discovered_adapters: list[tuple[str | None, str, str, str | None]], lora_substitute_id: int, base_reset: bool = False, + classifier_substitute_id: int | None = None, ) -> list[int]: """Build the token-exchange substitute ids, index-aligned with the control tokens. @@ -849,12 +850,14 @@ def build_substitute_token_ids( * base-reset: ``lora_substitute_id`` as well. The token is placed at a turn boundary, where the role-open marker is exactly what the decoder expects at that position. + * classifier: ``classifier_substitute_id``, the base's role-open marker. Args: discovered_adapters: ``(adapter_path, adapter_name, technology, source)`` tuples. lora_substitute_id: Probed sequence-start token id. base_reset: When True, prepend the base-reset slot's substitute so the list stays aligned with ``add_control_tokens(base_reset=True)``. + classifier_substitute_id: The base's role-open marker id. Returns: One substitute id per control token, in the same order. @@ -863,6 +866,10 @@ def build_substitute_token_ids( for adapter_path, _name, technology, _source in discovered_adapters: if technology == "alora": substitute_ids.append(get_alora_first_invocation_token_id(adapter_path)) + elif technology == "classifier": + if classifier_substitute_id is None: + raise ValueError("A classifier slot needs classifier_substitute_id.") + substitute_ids.append(classifier_substitute_id) else: substitute_ids.append(lora_substitute_id) return substitute_ids diff --git a/tests/composer/test_base_reset_token.py b/tests/composer/test_base_reset_token.py index a0010b6d..3869739c 100644 --- a/tests/composer/test_base_reset_token.py +++ b/tests/composer/test_base_reset_token.py @@ -22,6 +22,9 @@ _PROBE = ( "granite_switch.composer.compose_granite_switch._probe_lora_substitute_token_id" ) +_CLASSIFIER_SUB = ( + "granite_switch.composer.compose_granite_switch._classifier_substitute_token_id" +) _ALORA = "granite_switch.composer.tokenizer_setup.get_alora_first_invocation_token_id" @@ -50,7 +53,11 @@ class TestBuildControlTokenLists: def _run(self, base_reset): tokenizer = MockTokenizer(initial_vocab_size=500) - with patch(_PROBE, return_value=42), patch(_ALORA, return_value=77): + with ( + patch(_PROBE, return_value=42), + patch(_CLASSIFIER_SUB, return_value=43), + patch(_ALORA, return_value=77), + ): return build_control_token_lists(tokenizer, self._ADAPTERS, base_reset) def test_base_reset_grows_both_lists_and_leads_with_base(self): @@ -84,7 +91,11 @@ def test_compose_layout_is_the_one_the_engine_reads_as_base_reset(self): from .test_tokenizer_setup import MockTokenizer as _Tok tokenizer = _Tok(initial_vocab_size=500) - with patch(_PROBE, return_value=42), patch(_ALORA, return_value=77): + with ( + patch(_PROBE, return_value=42), + patch(_CLASSIFIER_SUB, return_value=43), + patch(_ALORA, return_value=77), + ): token_ids, _special, _subs = build_control_token_lists( tokenizer, self._ADAPTERS, base_reset=True ) @@ -125,7 +136,11 @@ def test_default_layout_keeps_the_adapter_offset(self): from .test_tokenizer_setup import MockTokenizer as _Tok tokenizer = _Tok(initial_vocab_size=500) - with patch(_PROBE, return_value=42), patch(_ALORA, return_value=77): + with ( + patch(_PROBE, return_value=42), + patch(_CLASSIFIER_SUB, return_value=43), + patch(_ALORA, return_value=77), + ): token_ids, _special, _subs = build_control_token_lists( tokenizer, self._ADAPTERS, base_reset=False ) diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py index b7482786..c3999d36 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_classifier_compose.py @@ -484,3 +484,25 @@ def test_classifier_marker_is_last_token(base): assert ids[-1] == control_id assert ids.count(control_id) == 1 assert tokenizer.decode(ids[:-1]) == plain + + +@pytest.mark.parametrize( + "base, role_open", + [ + ("ibm-granite/granite-4.0-micro", "<|start_of_role|>"), + ("ibm-granite/granite-4.1-3b", "<|start_of_role|>"), + ("ibm-granite/granite-4.2-3b", "<|im_start|>"), + ], +) +def test_classifier_substitute_is_the_bases_role_open_marker(base, role_open): + from granite_switch.composer.compose_granite_switch import ( + build_control_token_lists, + ) + + tokenizer = _tokenizer(base) + discovered = [(None, "safety", "classifier", None)] + _, _, substitutes = build_control_token_lists( + tokenizer, discovered, base_reset=False + ) + + assert substitutes == [tokenizer.convert_tokens_to_ids(role_open)] diff --git a/tests/composer/test_tokenizer_setup.py b/tests/composer/test_tokenizer_setup.py index d072dd91..2172ff80 100644 --- a/tests/composer/test_tokenizer_setup.py +++ b/tests/composer/test_tokenizer_setup.py @@ -344,6 +344,21 @@ def test_builtin_technology_uses_the_probed_substitute(self): assert subs == [42] + def test_classifier_uses_its_own_substitute(self): + adapters = [("/b", "code", "lora", None), (None, "safety", "classifier", None)] + + subs = build_substitute_token_ids( + adapters, lora_substitute_id=42, classifier_substitute_id=9 + ) + + assert subs == [42, 9] + + def test_classifier_without_its_substitute_raises(self): + adapters = [(None, "safety", "classifier", None)] + + with pytest.raises(ValueError, match="classifier_substitute_id"): + build_substitute_token_ids(adapters, lora_substitute_id=42) + class TestConfigureChatTemplate: """Structural tests for configure_chat_template — verify template assembly.""" diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index ad2581c3..9617ccc7 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -50,7 +50,12 @@ size. That is the whole artifact — a single linear layer that maps a base-mode logit per label. Train the head however you like (a frozen base model plus a trainable linear layer over the last -hidden state is the usual recipe). Here is a helper that writes a **synthetic** head. +hidden state is the usual recipe). At serving time the head reads the final hidden state of the +control token, which the model embeds as the base's role-open marker: `<|start_of_role|>` on +Granite 4.0/4.1, `<|im_start|>` on 4.2. So train on the chat-rendered prompt with no generation +prompt, with that marker appended, and read the hidden state there. + +Here is a helper that writes a **synthetic** head. ```python from pathlib import Path From 099b477dda890315fb8c28c3d2c559408b3628a2 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Sat, 10 Oct 2026 00:48:35 +0000 Subject: [PATCH 09/20] feat(classifier): support configurable classifier read layers Add per-slot classifier read-layer configuration to the manifest, composer, and model config. Score intermediate-layer classifiers as decoder layers run, in both the HF and vLLM backends, while preserving the active LoRA selection at the classifier marker. In vLLM, keep classifier scores aligned with sampled rows through the hidden-state/logits path. Add coverage for CUDA graph replay, chunked prefill, mixed classifier and plain requests, LoRA coactivation, and intermediate-layer reads. Intermediate reads require pipeline parallelism of 1 and are not supported with Shadow Residual. Speculative decoding and microbatching are unsupported with classifier slots, and classifier slot ids are limited to 256. Classifier markers must be the request's final token; vLLM misplaced-marker validation is deferred. Signed-off-by: IRENEKO --- .../composer/adapter_discovery.py | 12 +- .../composer/compose_granite_switch.py | 16 +- src/granite_switch/composer/compose_utils.py | 5 +- src/granite_switch/config.py | 72 ++- .../hf/modeling_granite_switch.py | 171 ++++--- src/granite_switch/hf/switch/multi.py | 43 +- src/granite_switch/token_exchange.py | 53 +- .../vllm/granite_switch_model.py | 300 +++++++---- src/granite_switch/vllm/switch/multi.py | 45 +- tests/composer/test_classifier_compose.py | 14 +- tests/hf/test_classifier.py | 124 ++++- tests/unit/test_config.py | 69 +++ tests/unit/test_token_exchange.py | 43 ++ .../_classifier_chunked_prefill_worker.py | 484 ++++++++++++++++++ tests/vllm/_model_forward_tests.py | 48 +- tests/vllm/test_classifier_capture_safety.py | 194 +++++++ tests/vllm/test_classifier_chunked_prefill.py | 335 +++++++++++- tests/vllm/test_classifier_read_points.py | 93 ---- tutorials/guides/classifier_slots.md | 50 +- 19 files changed, 1819 insertions(+), 352 deletions(-) create mode 100644 tests/vllm/test_classifier_capture_safety.py delete mode 100644 tests/vllm/test_classifier_read_points.py diff --git a/src/granite_switch/composer/adapter_discovery.py b/src/granite_switch/composer/adapter_discovery.py index 8b6e099f..eb5b334e 100644 --- a/src/granite_switch/composer/adapter_discovery.py +++ b/src/granite_switch/composer/adapter_discovery.py @@ -112,7 +112,7 @@ def discover_adapters_from_yaml( """Discover adapters from a YAML manifest file. Reads a YAML manifest that maps adapter names to their paths and types. - An entry can also mark itself a **classifier** slot with two extra keys: + An entry can also mark itself a **classifier** slot with these extra keys: .. code-block:: yaml @@ -121,6 +121,7 @@ def discover_adapters_from_yaml( type: lora # lora/alora technology (token placement) kind: classifier # default "lora" when omitted labels: [faithful, hallucinated] + layer: 12 # optional decoder layer; omit for the final layer ``type`` and ``kind`` are independent axes: ``type`` is the LoRA technology (lora vs alora, which drives token placement), ``kind`` is whether the slot @@ -128,6 +129,9 @@ def discover_adapters_from_yaml( without them an entry is a plain LoRA adapter. The label words become token ids at compose time (see :func:`tokenizer_setup.resolve_label_token_ids`). + ``layer: K`` reads the residual stream after zero-based decoder layer ``K`` + (switch cache layers excluded). Omitting it reads the final post-norm state. + Args: manifest_path: Path to the YAML manifest file. @@ -136,7 +140,7 @@ def discover_adapters_from_yaml( List of ``(adapter_path, adapter_name, technology, source)`` tuples. The source is set to the manifest path for traceability. ``classifier_meta`` maps name -> - ``{"kind", "labels"}`` for each entry that set ``kind``/``labels``; + ``{"kind", "labels", "layer"}`` for each entry that set any of them; names not in it are plain LoRA slots. """ import yaml @@ -160,10 +164,12 @@ def discover_adapters_from_yaml( # Optional classifier slots. kind = info.get("kind") labels = info.get("labels") - if kind is not None or labels is not None: + layer = info.get("layer") + if kind is not None or labels is not None or layer is not None: classifier_meta[name] = { "kind": kind or "lora", "labels": labels, + "layer": layer, } return found, classifier_meta diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index d620b93b..0882b044 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -1281,10 +1281,13 @@ def build(): # counts and the padded bank width from this. per_slot_label_names: list[list[str] | None] = [] per_slot_label_token_ids: list[list[int] | None] = [] + # Full width too: each classifier slot's read layer (None = final). + per_slot_read_layer: list[int | None] = [] for name, kind in zip(adapter_names, adapter_kinds): if kind != "classifier": per_slot_label_names.append(None) per_slot_label_token_ids.append(None) + per_slot_read_layer.append(None) continue labels = classifier_meta[name].get("labels") if not labels: @@ -1300,14 +1303,21 @@ def build(): resolve_label_token_ids(tokenizer, label_names) ) + per_slot_read_layer.append(classifier_meta[name].get("layer")) + optional_kwargs["adapter_kinds"] = adapter_kinds optional_kwargs["classifier_label_token_ids"] = per_slot_label_token_ids + optional_kwargs["classifier_read_layers"] = per_slot_read_layer print("\nClassifier slots:") - for name, words, ids in zip( - adapter_names, per_slot_label_names, per_slot_label_token_ids + for name, words, ids, layer in zip( + adapter_names, + per_slot_label_names, + per_slot_label_token_ids, + per_slot_read_layer, ): if ids is not None: - print(f" {name} → labels {words} → token ids {ids}") + where = "final norm" if layer is None else f"layer {layer}" + print(f" {name} → labels {words} → token ids {ids} @ {where}") model = GraniteSwitchComposer.from_base_and_adapters( base_model_name_or_path=base_model_local_path, diff --git a/src/granite_switch/composer/compose_utils.py b/src/granite_switch/composer/compose_utils.py index b8982a1d..32e90913 100644 --- a/src/granite_switch/composer/compose_utils.py +++ b/src/granite_switch/composer/compose_utils.py @@ -175,8 +175,9 @@ def from_base_and_adapters( # adapter_ranks/adapter_alphas stay full external width (one entry per # slot, classifiers included) so config.adapter_ranks is length - # num_adapters. A classifier slot's LoRA row never fires: split_indices - # zeroes the LoRA stream at classifier positions. + # num_adapters. A classifier slot itself has no LoRA weights. A + # separately selected LoRA slot can remain active through a classifier + # marker because classifier selection does not overwrite adapter state. if lora_adapter_paths: lora_rank, lora_alpha, lora_ranks, lora_alphas = detect_lora_config( lora_adapter_paths diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index 37e795c2..0debdb05 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -27,12 +27,13 @@ class GraniteSwitchConfig(GraniteMoeHybridConfig): substitution. Args: - num_adapters (int): Number of LoRA adapters available. Default: 0 (no adapters). - This counts real LoRA adapters only (not base). Index 0 always means "base / no adapter". + num_adapters (int): Number of adapter slots available (LoRA and + classifier). Default: 0. Index 0 always means "base / no adapter". adapter_token_ids (List[int]): Token IDs for adapter control. - Length: num_adapters (one token per real adapter). Must be unique. - adapter_token_ids[i] activates adapter i+1 (1-indexed output). - Output 0 = base (implicit default, no token needed to return to base). + Length: num_adapters (or num_adapters + 1 with a leading base-reset + token). Must be unique. A LoRA control token selects that adapter; + a classifier control token selects its head without changing the + active adapter. Output 0 = base. adapter_substitute_token_ids (List[int]): Token IDs whose embeddings replace the control-token embeddings before the decoder runs. Length: num_adapters (or num_adapters + 1 with a leading base-reset @@ -184,6 +185,7 @@ def __init__( # Classifier substitution parameters adapter_kinds: list[str] | None = None, classifier_label_token_ids: list[list[int] | None] | None = None, + classifier_read_layers: list[int | None] | None = None, # vLLM residual-norm convention (for bit-exact skinning equivalence) fused_add_norm: bool = False, # Parent class defaults (Granite 4 dense configuration) @@ -385,10 +387,9 @@ def __init__( self.adapter_ranks = adapter_ranks # Each adapter slot is either a LoRA adapter ("lora", the default) or - # a classifier head ("classifier"). A classifier slot still uses a - # control token and an adapter index; the switch routes its token to - # that index like lora adapter, but the index is tagged as a classifier - # so the LoRA path no-ops there and the classifier head runs instead. + # a classifier head ("classifier"). Classifier markers are orthogonal + # probes: they select a head without replacing the currently active + # adapter selection. The probe slot is looked up from the marker token. _VALID_KINDS = ("lora", "classifier") if adapter_kinds is None: adapter_kinds = ["lora"] * num_adapters @@ -405,7 +406,6 @@ def __init__( f"got invalid entries {invalid}" ) self.adapter_kinds = adapter_kinds - # Classifier metadata resolved at compose time. if classifier_label_token_ids is not None: if len(classifier_label_token_ids) != num_adapters: @@ -441,8 +441,53 @@ def __init__( "verdict as a label word, so the per-label token ids are required. " "Re-compose so config.json carries classifier_label_token_ids." ) + # vLLM packs each token's classifier slot id into the bf16 hidden states, + # which hold integers exactly only up to 256. + if any( + kind == "classifier" and slot > 256 + for slot, kind in enumerate(adapter_kinds, start=1) + ): + raise ValueError( + "classifier slot ids must be <= 256; classifier slots are numbered " + "after LoRA slots, so compose at most 256 adapters." + ) self.classifier_label_token_ids = classifier_label_token_ids + # Per slot: zero-based decoder layer a classifier reads (switch cache + # layers excluded); None reads the final post-norm hidden state. + if classifier_read_layers is not None: + if len(classifier_read_layers) != num_adapters: + raise ValueError( + f"classifier_read_layers length ({len(classifier_read_layers)}) " + f"must equal num_adapters ({num_adapters})." + ) + num_decoder_layers = self.num_hidden_layers - SWITCH_CACHE_LAYERS + for i, (kind, layer) in enumerate( + zip(adapter_kinds, classifier_read_layers) + ): + if layer is None: + continue + if kind != "classifier": + raise ValueError( + f"non-classifier slot {i} must have no read layer; got {layer!r}." + ) + # bool is an int subclass. + if ( + not isinstance(layer, int) + or isinstance(layer, bool) + or not 0 <= layer < num_decoder_layers + ): + raise ValueError( + f"classifier_read_layers[{i}] = {layer!r} must be an int in " + f"0..{num_decoder_layers - 1}; omit it for the final layer." + ) + if self.dual_stream: + raise ValueError( + f"classifier slot {i} reads layer {layer}, but intermediate " + f"layers are not supported with dual_stream=True." + ) + self.classifier_read_layers = classifier_read_layers + # Derived: per-slot label counts (0 on LoRA slots) and the padded bank # width the classifier head is built with. self.classifier_num_labels_per_slot = [ @@ -453,10 +498,11 @@ def __init__( # Derived: the control-token ids that mark a classifier slot. Both backends # locate a verdict's read point by matching these in the original input_ids, # so the id set is defined once here. + classifier_token_offset = len(adapter_token_ids or []) - num_adapters self.classifier_control_token_ids = [ - tid - for tid, kind in zip(adapter_token_ids or [], adapter_kinds) - if kind == "classifier" + (adapter_token_ids or [])[i + classifier_token_offset] + for i, kind in enumerate(adapter_kinds) + if kind == "classifier" and adapter_token_ids is not None ] # Default LoRA target module groups. diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index 3bbd013e..0729b26a 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -419,12 +419,10 @@ def __init__(self, config: GraniteSwitchConfig): self.switch = None self.adapter_token_ids = None - # Classifier heads: optional alternative for LoRA adapters. A classifier - # slot shares the adapter index/token machinery, but its index is tagged - # as "classifier" (via the switch's adapter_kind_lut). The switch's - # split_indices() routes those positions out of the LoRA stream (so LoRA - # no-ops there) and into a classifier stream, which this head consumes - # from the final post-norm hidden state. + # Classifier heads share the adapter control-token machinery. The marker + # selects a classifier head independently; the active adapter selection + # remains in effect through that token. Intermediate read layers are + # scored as the decoder runs, retaining only marker-row verdicts. if "classifier" in getattr(config, "adapter_kinds", []): self.classifier_head = SwitchedClassifierHead( hidden_size=config.hidden_size, @@ -604,12 +602,17 @@ def forward( device=device, ) - # Split the single adapter_indices stream into a LoRA stream and - # a classifier stream. Classifier positions are zeroed in lora_indices - # (so the LoRA path no-ops there). + # Keep persistent adapter routing separate from the classifier probe: + # classifier markers select a head from their token id and do not + # replace the currently active adapter. if self.switch is not None: + classifier_probe_indices = ( + self.switch.classifier_indices_from_tokens(input_ids) + if input_ids is not None + else None + ) lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices + adapter_indices, classifier_probe_indices ) else: lora_indices = adapter_indices @@ -622,10 +625,10 @@ def forward( inputs_embeds = self.embed_tokens(modified_input_ids) inputs_embeds = inputs_embeds * self.embedding_multiplier - # Expose adapter_indices for tests and debugging. _last_adapter_indices is - # the full stream (LoRA + classifier); _last_classifier_indices is the - # classifier-only stream the verdict reads. + # Expose routing for tests/debugging: adapter selection is persistent, + # while classifier indices come independently from marker token ids. self._last_adapter_indices = adapter_indices + self._last_lora_indices = lora_indices self._last_classifier_indices = classifier_indices position_embeddings = None @@ -646,7 +649,79 @@ def forward( all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None - for decoder_layer in self.layers: + # Resolve classifier read positions before the decoder runs. This lets + # intermediate heads consume only their marker rows as each layer + # produces them, without retaining full sequence activations per layer. + self._last_classifier_verdict = None + classifier_rows = classifier_pos = classifier_slots = None + classifier_verdict = None + classifier_markers_by_layer = {} + classifier_final_markers = None + if self.classifier_head is not None: + if input_ids is None: + if bool((classifier_indices > 0).any()): + raise RuntimeError( + "Classifier requests need input_ids, not inputs_embeds." + ) + else: + marker_ids = torch.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + is_marker = torch.isin(input_ids, marker_ids) + if bool(is_marker.any()): + classifier_rows, classifier_pos = is_marker.nonzero(as_tuple=True) + classifier_slots = classifier_indices[ + classifier_rows, classifier_pos + ] + + # Each marker must be the last real token of its request, + # since the verdict replaces that position's logits. + if attention_mask is not None and attention_mask.dim() == 2: + real = attention_mask[:, -seq_length:].long() + cols = torch.arange(seq_length, device=real.device) + last = (real * cols).max(dim=1).values + else: + last = torch.full( + (batch_size,), seq_length - 1, device=input_ids.device + ) + misplaced = classifier_pos != last[classifier_rows] + if bool(misplaced.any()): + bad = sorted(set(classifier_rows[misplaced].tolist())) + raise RuntimeError( + f"Classifier request(s) {bad}: control token is not " + "their last token." + ) + + read_layers = self.config.classifier_read_layers + if read_layers is None: + read_layers = [None] * self.config.num_adapters + layer_to_marker_indices = {} + final_marker_indices = [] + for marker_idx, slot in enumerate(classifier_slots.tolist()): + layer = read_layers[slot - 1] if slot > 0 else None + if layer is None: + final_marker_indices.append(marker_idx) + else: + layer_to_marker_indices.setdefault(layer, []).append( + marker_idx + ) + classifier_markers_by_layer = { + layer: torch.tensor(indices, device=input_ids.device) + for layer, indices in layer_to_marker_indices.items() + } + classifier_final_markers = torch.tensor( + final_marker_indices, device=input_ids.device + ) + classifier_verdict = inputs_embeds.new_zeros( + ( + classifier_rows.numel(), + self.config.max_classifier_labels, + ) + ) + + for layer_idx, decoder_layer in enumerate(self.layers): if output_hidden_states: all_hidden_states += (hidden_states,) @@ -668,6 +743,15 @@ def forward( else: hidden_states = layer_outputs[0] + marker_indices = classifier_markers_by_layer.get(layer_idx) + if marker_indices is not None: + marker_hidden = hidden_states[ + classifier_rows[marker_indices], classifier_pos[marker_indices] + ] + classifier_verdict[marker_indices] = self.classifier_head( + marker_hidden, classifier_slots[marker_indices] + ) + if output_attentions: if layer_outputs[1] is not None: all_self_attns += (layer_outputs[1],) @@ -677,51 +761,22 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - # Classifier heads read the final post-norm hidden state (the same - # representation the LM head consumes), only at classifier markers. Each - # marker must be its row's last real token, where the exit rewrites that - # row's logits into a label word. - self._last_classifier_verdict = None - if self.classifier_head is not None: - if input_ids is None: - if bool((classifier_indices > 0).any()): - raise RuntimeError( - "Classifier requests need input_ids, not inputs_embeds." - ) - else: - marker_ids = torch.tensor( - self.config.classifier_control_token_ids, - device=input_ids.device, - dtype=input_ids.dtype, + if classifier_rows is not None: + if classifier_final_markers.numel() > 0: + final_hidden = hidden_states[ + classifier_rows[classifier_final_markers], + classifier_pos[classifier_final_markers], + ] + classifier_verdict[classifier_final_markers] = self.classifier_head( + final_hidden, classifier_slots[classifier_final_markers] ) - is_marker = torch.isin(input_ids, marker_ids) # [batch, seq_len] - if bool(is_marker.any()): - rows, pos = is_marker.nonzero(as_tuple=True) - # Last real token of each row among this step's tokens. - if attention_mask is not None and attention_mask.dim() == 2: - real = attention_mask[:, -seq_length:].long() - cols = torch.arange(seq_length, device=real.device) - last = (real * cols).max(dim=1).values - else: - last = torch.full( - (batch_size,), seq_length - 1, device=input_ids.device - ) - misplaced = pos != last[rows] - if bool(misplaced.any()): - bad = sorted(set(rows[misplaced].tolist())) - raise RuntimeError( - f"Classifier request(s) {bad}: control token is not " - "their last token." - ) - slots = classifier_indices[rows, pos] - verdict = self.classifier_head(hidden_states[rows, pos], slots) - self._last_classifier_verdict = ( - rows, - pos, - verdict, - slots, - seq_length, - ) + self._last_classifier_verdict = ( + classifier_rows, + classifier_pos, + classifier_verdict, + classifier_slots, + seq_length, + ) if not return_dict: return tuple( diff --git a/src/granite_switch/hf/switch/multi.py b/src/granite_switch/hf/switch/multi.py index f98e5773..549ef415 100644 --- a/src/granite_switch/hf/switch/multi.py +++ b/src/granite_switch/hf/switch/multi.py @@ -69,6 +69,7 @@ from ...token_exchange import ( apply_token_exchange, build_adapter_kind_lut, + build_classifier_control_luts, build_control_to_substitute_lut, split_adapter_indices, ) @@ -294,6 +295,13 @@ def __init__( self.register_buffer("control_to_substitute_lut", lut, persistent=True) else: self.control_to_substitute_lut = None + marker_lut, slot_lut = build_classifier_control_luts(config) + if marker_lut is not None: + self.register_buffer("classifier_control_lut", marker_lut, persistent=False) + self.register_buffer("classifier_slot_lut", slot_lut, persistent=False) + else: + self.classifier_control_lut = None + self.classifier_slot_lut = None def _apply(self, *args, **kwargs): """Keep ``codebook`` in fp32 across ``.to()`` / ``.half()`` / ``.bfloat16()``. @@ -323,13 +331,26 @@ def _apply(self, *args, **kwargs): return out def split_indices( - self, adapter_indices: torch.Tensor + self, + adapter_indices: torch.Tensor, + classifier_probe_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Split ``[batch, seq_len]`` or ``[num_tokens]`` into (lora, classifier). See :func:`granite_switch.token_exchange.split_adapter_indices`. """ - return split_adapter_indices(self.adapter_kind_lut, adapter_indices) + return split_adapter_indices( + self.adapter_kind_lut, + adapter_indices, + classifier_probe_indices, + ) + + def classifier_indices_from_tokens( + self, input_ids: torch.Tensor + ) -> torch.Tensor | None: + if self.classifier_slot_lut is None: + return None + return self.classifier_slot_lut[input_ids] @property def num_cache_layers(self) -> int: @@ -384,9 +405,15 @@ def forward( # [B, S, 1] == [1, 1, A] -> [B, S, A] matches = input_ids.unsqueeze(2) == adapter_token_ids.unsqueeze(0).unsqueeze(0) is_control_token = matches.any(dim=2) # [B, S] + route_control_token = is_control_token + if self.classifier_control_lut is not None: + # A classifier marker probes the current residual stream; it must + # not overwrite the persistent adapter selection in coded memory. + is_classifier_marker = self.classifier_control_lut[input_ids] + route_control_token = is_control_token & ~is_classifier_marker # expert_id = argmax + offset (offset selects the layout; see __init__). expert_ids = torch.where( - is_control_token, + route_control_token, matches.long().argmax(dim=2) + self._expert_id_offset, torch.zeros_like(input_ids, dtype=torch.long), ) # [B, S] @@ -401,8 +428,10 @@ def forward( device=device, dtype=torch.float32, ) - # Un-mask (0) dim 0 at anchor + control tokens (arithmetic where). - anchor_or_control = is_counting_anchor | is_control_token # [B, S] + # Un-mask at the anchor and adapter-selection controls. Classifier + # markers are excluded because they do not change persistent adapter + # selection. + anchor_or_control = is_counting_anchor | route_control_token # [B, S] key_states_count[:, 0, :, 0] = torch.where( anchor_or_control, torch.zeros_like(key_states_count[:, 0, :, 0]), @@ -528,7 +557,7 @@ def _causal_mask(kv_len: int) -> torch.Tensor: bsz, q_len, self.memory_dim ) # [B, S, memory_dim], fp32 - write_mask = is_control_token.unsqueeze(2).to(torch.float32) # [B, S, 1] + write_mask = route_control_token.unsqueeze(2).to(torch.float32) # [B, S, 1] key_states_memory = torch.zeros( (bsz, 1, q_len, self.memory_head_dim), @@ -547,7 +576,7 @@ def _causal_mask(kv_len: int) -> torch.Tensor: ) value_states_memory[:, 0, :, 0] = expert_ids.to( torch.float32 - ) * is_control_token.to(torch.float32) + ) * route_control_token.to(torch.float32) if past_key_values is not None: cache_kwargs = {"cache_position": cache_position} diff --git a/src/granite_switch/token_exchange.py b/src/granite_switch/token_exchange.py index e0a37dff..d73d1aa6 100644 --- a/src/granite_switch/token_exchange.py +++ b/src/granite_switch/token_exchange.py @@ -161,9 +161,9 @@ def build_adapter_kind_lut(config) -> torch.Tensor | None: ``[num_adapters + 1]``, indexed by ADAPTER INDEX: entry ``i`` is True when adapter index ``i`` is a classifier slot (index 0, base, is always False). - ``None`` when no slot is a classifier, which lets the switch skip the split - entirely. Built from ``config.adapter_kinds``, so both backends' switches - agree on the table by construction. + ``None`` when no slot is a classifier. This supports the legacy split when + no token-derived classifier probe indices are supplied. Built from + ``config.adapter_kinds``, so both backends agree on the table. """ kinds = getattr(config, "adapter_kinds", None) if config is not None else None if kinds is None or not any(k == "classifier" for k in kinds): @@ -175,21 +175,60 @@ def build_adapter_kind_lut(config) -> torch.Tensor | None: return lut +def build_classifier_control_luts( + config, +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Build token-id LUTs for classifier markers. + + The first LUT marks classifier control tokens; the second returns the + classifier head slot (1-based) for each token. These LUTs are separate + from ``adapter_kind_lut`` because an orthogonal classifier probe must not + replace the persistent LoRA adapter selected earlier in the request. + """ + kinds = getattr(config, "adapter_kinds", None) + control_ids = getattr(config, "adapter_token_ids", None) + if not kinds or not control_ids or not any(k == "classifier" for k in kinds): + return None, None + offset = len(control_ids) - len(kinds) + if offset not in (0, 1): + raise ValueError( + "adapter_token_ids must have one id per adapter, optionally with " + "a leading base-reset token" + ) + vocab_size = max( + int(getattr(config, "vocab_size", 0)), + max(control_ids) + 1, + ) + is_marker = torch.zeros(vocab_size, dtype=torch.bool) + slot_lut = torch.zeros(vocab_size, dtype=torch.long) + for slot, kind in enumerate(kinds, start=1): + if kind != "classifier": + continue + token_id = control_ids[slot - 1 + offset] + is_marker[token_id] = True + slot_lut[token_id] = slot + return is_marker, slot_lut + + def split_adapter_indices( - kind_lut: torch.Tensor | None, adapter_indices: torch.Tensor + kind_lut: torch.Tensor | None, + adapter_indices: torch.Tensor, + classifier_probe_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Split a per-token index stream into LoRA and classifier streams. Shapes are preserved (``[batch, seq_len]`` on HF, ``[total_tokens]`` on vLLM), so this serves both backends: - - ``lora_indices``: classifier positions zeroed, so the LoRA path no-ops - there. - - ``classifier_indices``: the adapter index at classifier positions, else 0. + - When classifier marker indices are provided, they select the classifier + head independently and the active adapter stream is preserved. + - Without marker indices, the legacy kind LUT split is used. With ``kind_lut`` None (no classifier slots) the LoRA stream passes through and the classifier stream is all zeros. """ + if classifier_probe_indices is not None: + return adapter_indices, classifier_probe_indices if kind_lut is None: return adapter_indices, torch.zeros_like(adapter_indices) is_classifier = kind_lut[adapter_indices] diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index 51d3c456..82d6c618 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -168,12 +168,9 @@ def __init__( self.lora_meta = None self.lora_ctx = None - # Classifier heads: optional alternative for LoRA adapters. A classifier - # slot shares the adapter index/token machinery, but its index is tagged - # as "classifier" (via the switch's adapter_kind_lut). The switch's - # split_indices() routes those positions out of the LoRA stream (so LoRA - # no-ops there) and into a classifier stream, which this head consumes - # from the final post-norm hidden state. + # Classifier markers share the control-token machinery but select the + # classifier head independently, preserving the active adapter. The head + # scores each slot at its read layer (the final post-norm state by default). if "classifier" in getattr(config, "adapter_kinds", []): from .core import SwitchedClassifierHead @@ -182,8 +179,63 @@ def __init__( num_classifier_slots=config.num_adapters, max_num_labels=config.max_classifier_labels, ) + + # Packed verdicts assume the ordinary sampler's row selection. + if vllm_config.speculative_config is not None: + raise ValueError( + "Classifier slots do not support speculative decoding." + ) + if vllm_config.parallel_config.use_ubatching: + raise ValueError("Classifier slots do not support microbatching.") + + # Built once: copying a list to CUDA inside forward breaks capture. + self.register_buffer( + "classifier_control_token_ids_t", + torch.tensor(config.classifier_control_token_ids, dtype=torch.long), + persistent=False, + ) + + is_classifier = [kind == "classifier" for kind in config.adapter_kinds] + # Index 0 is a sentinel so a 1-based slot id indexes it directly. + self.register_buffer( + "classifier_slot_mask_t", + torch.tensor([False, *is_classifier]), + persistent=False, + ) + + # One buffer of slot ids per read layer (None = final). Store buffer + # names, not tensors: Module.to() replaces registered buffers. + read_layers = config.classifier_read_layers or [None] * num_adapters + slots_by_layer: dict[int | None, list[int]] = {} + for slot, (is_clf, layer) in enumerate( + zip(is_classifier, read_layers), start=1 + ): + if is_clf: + slots_by_layer.setdefault(layer, []).append(slot) + self._classifier_final_layer_group = None + self._classifier_layer_groups: dict[int, str] = {} + for layer, slot_ids in slots_by_layer.items(): + name = f"_classifier_slots_{'final' if layer is None else layer}" + self.register_buffer( + name, torch.tensor(slot_ids, dtype=torch.long), persistent=False + ) + if layer is None: + self._classifier_final_layer_group = name + else: + self._classifier_layer_groups[layer] = name + # Intermediate scores are not carried across pipeline ranks. + if ( + self._classifier_layer_groups + and vllm_config.parallel_config.pipeline_parallel_size > 1 + ): + raise ValueError( + "Intermediate-layer classifier slots do not support " + "pipeline parallelism." + ) else: self.classifier_head = None + self._classifier_final_layer_group = None + self._classifier_layer_groups = {} # 3. Base transformer layers with custom LoRA # @@ -334,30 +386,34 @@ def forward( ) modified_input_ids = input_ids - # Split the single adapter_indices stream into a LoRA stream and a - # classifier stream. Classifier positions are zeroed in lora_indices - # (so the LoRA path no-ops there); the classifier head reads them on - # the last rank via classifier_indices. + # Split routing state from classifier probe state. In the default + # mode classifier slots are removed from the LoRA stream; markers + # get a separate probe index while the previously selected adapter + # remains active at that token. if self.switch is not None: + classifier_probe_indices = ( + self.switch.classifier_indices_from_tokens(input_ids) + if input_ids is not None + else None + ) lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices + adapter_indices, classifier_probe_indices ) else: lora_indices = adapter_indices classifier_indices = torch.zeros_like(adapter_indices) - # Prepare kernel metadata ONCE for all decoder layers - # (adaptation-specific). The LoRA stream (classifier positions zeroed) - # is what the fused kernel routes on, so classifier tokens no-op the - # LoRA path. + # Prepare kernel metadata ONCE for all decoder layers. The fused + # kernels route on lora_indices, including the marker token when + # coactivation is enabled. if self.lora_meta is not None and self.lora_ctx is not None: self.decoder_interface.prepare_kernel_meta( self.lora_meta, lora_indices, self.lora_ctx ) - # Store metadata in intermediate_tensors for pipeline parallelism. - # Only adapter_indices is propagated; classifier_indices is derived - # from it on the last rank via split_indices (no new PP tensor). + # Store the persistent adapter state for pipeline parallelism. + # Classifier probe indices are reconstructed from this rank's + # original input_ids and the config-derived marker LUT. if intermediate_tensors is None: intermediate_tensors = IntermediateTensors({}) intermediate_tensors["adapter_indices"] = adapter_indices @@ -366,19 +422,25 @@ def forward( # token-leading adapter_indices received through PP. if intermediate_tensors is not None: adapter_indices = intermediate_tensors["adapter_indices"] - # Recompute the split from the PP-propagated stream so this - # rank's Punica metadata sees classifier positions as base. + # Recompute both streams from the PP-propagated adapter state + # and this rank's token ids. The classifier marker LUT is a + # non-persistent config-derived buffer on every rank. if self.switch is not None: + classifier_probe_indices = ( + self.switch.classifier_indices_from_tokens(input_ids) + if input_ids is not None + else None + ) lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices + adapter_indices, classifier_probe_indices ) else: lora_indices = adapter_indices classifier_indices = torch.zeros_like(adapter_indices) if self.lora_ctx is not None: - # Route on the LoRA stream (classifier positions zeroed) so - # classifier tokens no-op the LoRA path on this PP rank too. + # Reuse this rank's LoRA stream, including the marker token + # when coactivation is enabled. self.decoder_interface.prepare_kernel_meta( self.lora_meta, lora_indices, self.lora_ctx ) @@ -424,9 +486,27 @@ def forward( # threads (hidden_states, residual); SR re-stacks to [2M,H] rank-locally # in enter_decoder_stack and collapses it in exit_decoder_stack. All # per-forward metadata (kernel meta + hiding masks) is on the shared ctx. + # Classifier slots reading an intermediate layer are scored as soon as + # that layer runs; only the small score tensor is carried forward. + if self.classifier_head is not None: + classifier_scores = torch.zeros( + classifier_indices.shape[0], + self.config.max_classifier_labels, + device=classifier_indices.device, + dtype=hidden_states.dtype, + ) state = self.decoder_interface.enter_decoder_stack(hidden_states, residual) for i in range(self.start_layer, self.end_layer): state = self.decoder_interface.run_layer(self.layers[i], positions, state) + group = self._classifier_layer_groups.get(i) + if group is not None: + # Layer i's output before the final norm: residual + block output. + classifier_scores = self._score_group( + group, + state.residual + state.hidden_states, + classifier_indices, + classifier_scores, + ) if get_pp_group().is_last_rank: # Adaptation-specific finalize: LoRA folds the last residual via @@ -438,8 +518,16 @@ def forward( state, adapter_indices, self.norm, self.config ) - if self.lora_ctx is not None: - self.lora_ctx.classifier_indices = classifier_indices + # Pack scores and slot ids onto the hidden states so vLLM's own + # logits_indices gather selects them; compute_logits unpacks them. + if self.classifier_head is not None: + classifier_scores, classifier_slot = self._compute_classifier_outputs( + input_ids, hidden_states, classifier_indices, classifier_scores + ) + hidden_states = torch.cat( + [hidden_states, classifier_scores, classifier_slot.unsqueeze(-1)], + dim=-1, + ) return hidden_states else: @@ -453,52 +541,48 @@ def forward( intermediate_tensors["residual"] = r_out return intermediate_tensors - def _classifier_read_points( + def _score_group( self, + group: str, + hidden: torch.Tensor, classifier_indices: torch.Tensor, - hidden_states: torch.Tensor, + classifier_scores: torch.Tensor, + ) -> torch.Tensor: + """Write the scores of the slots in buffer ``group``, read from ``hidden``.""" + in_group = torch.isin(classifier_indices, getattr(self, group)) + slot_logits = self.classifier_head(hidden, classifier_indices) + return torch.where(in_group.unsqueeze(-1), slot_logits, classifier_scores) + + def _compute_classifier_outputs( + self, input_ids: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None: - """``(rows, hidden, slots, num_reqs)`` for this pass's classifier markers. + final_hidden_states: torch.Tensor, + classifier_indices: torch.Tensor, + classifier_scores: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Add final-layer scores and return per-token (scores, slot ids). - Each marker must be its request's last token, so it reports to that - request's logit row. ``None`` when the pass holds no marker. + Only marker tokens carry a nonzero slot id. Callers must place one + classifier marker last in each request; misplaced or duplicate markers + are not rejected (request-boundary validation is deferred). """ - if input_ids is None: - raise RuntimeError("Classifier requests need input_ids, not inputs_embeds.") - marker_ids = torch.tensor( - self.config.classifier_control_token_ids, - device=input_ids.device, - dtype=input_ids.dtype, + is_marker = ( + torch.isin(input_ids, self.classifier_control_token_ids_t) + if input_ids is not None + else torch.zeros_like(classifier_indices, dtype=torch.bool) ) - is_marker = torch.isin(input_ids, marker_ids) # [total_tokens] - if not bool(is_marker.any()): - return None - - from vllm.forward_context import get_forward_context - - attn_metadata = get_forward_context().attn_metadata - # v1 keys this by layer name (a dict), or a list of such dicts under - # microbatching. They all share query_start_loc, so grab any one. - if isinstance(attn_metadata, list): - attn_metadata = attn_metadata[0] if attn_metadata else None - if isinstance(attn_metadata, dict): - attn_metadata = next(iter(attn_metadata.values()), None) - query_start_loc = getattr(attn_metadata, "query_start_loc", None) - if query_start_loc is None: - return None - - pos = is_marker.nonzero(as_tuple=True)[0] # [num_markers] - ends = query_start_loc[1:].long() - # The request whose slice [start, end) holds each marker = its logit row. - rows = torch.searchsorted(ends, pos, right=True) - misplaced = pos != ends[rows] - 1 - if bool(misplaced.any()): - reqs = sorted(set(rows[misplaced].tolist())) - raise RuntimeError( - f"Classifier request(s) {reqs}: control token is not their last token." + # Exact in bf16: the config limits classifier slot ids to 256. + classifier_slot = torch.where( + is_marker, classifier_indices, torch.zeros_like(classifier_indices) + ).to(final_hidden_states.dtype) + if self._classifier_final_layer_group is not None: + classifier_scores = self._score_group( + self._classifier_final_layer_group, + final_hidden_states, + classifier_indices, + classifier_scores, ) - return rows, hidden_states[pos], classifier_indices[pos].long(), len(ends) + return classifier_scores, classifier_slot @MULTIMODAL_REGISTRY.register_processor( @@ -683,31 +767,13 @@ def forward( intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: - """Forward pass returning hidden states.""" - hidden_states = self.model( + """Forward pass returning hidden states (classifier columns packed on).""" + return self.model( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, ) - self._pending_verdict = None - classifier_head = getattr(self.model, "classifier_head", None) - if classifier_head is not None and isinstance(hidden_states, torch.Tensor): - classifier_indices = getattr( - self.model.lora_ctx, "classifier_indices", None - ) - if classifier_indices is not None: - read = self.model._classifier_read_points( - classifier_indices, hidden_states, input_ids - ) - if read is not None: - rows, marker_hidden, slots, num_reqs = read - verdict = classifier_head( - marker_hidden, slots - ) # [num_markers, max_labels] - self._pending_verdict = (rows, verdict, slots, num_reqs) - - return hidden_states def compute_logits( self, @@ -724,41 +790,59 @@ def compute_logits( have to act where adapter_indices and hidden_states still share the same token dimension. - Classifier exit: the classifier request reports its verdict as a generated - token, like a guardian LoRA. For each classifier request we rewrite its logit - row: set the whole vocab to -inf, then place the per-label scores on each - label word's token id (``classifier_label_token_ids``) so the sampler - emits that word. Non-classifier rows keep their normal LM logits. + Classifier exit: the classifier request reports its verdict as a + generated token, like a guardian LoRA. With a classifier head the + hidden states carry packed (scores, slot id) columns; split them off, + run the LM head on the real hidden state, then rewrite each classifier + row: the whole vocab to -inf and the per-label scores onto the label + words' token ids (``classifier_label_token_ids``). Other rows keep + their LM logits. """ - logits = self.logits_processor(self.lm_head, hidden_states) - return self._apply_classifier_verdict(logits) + if self.model.classifier_head is not None: + hidden_size = self.config.hidden_size + real_hidden = hidden_states[:, :hidden_size] + verdict = hidden_states[:, hidden_size:-1] + slot = hidden_states[:, -1] + else: + real_hidden = hidden_states + verdict = slot = None + logits = self.logits_processor(self.lm_head, real_hidden) + return self._apply_classifier_verdict(logits, verdict, slot) def _apply_classifier_verdict( self, logits: torch.Tensor | None, + verdict: torch.Tensor | None, + slot: torch.Tensor | None, ) -> torch.Tensor | None: - """Rewrite classifier requests' logit rows to emit their label word.""" - pending = getattr(self, "_pending_verdict", None) - if logits is None or pending is None: + """Rewrite classifier requests' logit rows to emit their label word. + + Rows are classified by the slot mask, not just ``slot > 0``: vLLM's + profiling sampler run feeds random hidden states, so the packed slot + column can land on a LoRA slot, which has no label ids. + """ + if logits is None or verdict is None or slot is None: + return logits + + slot_mask = self.model.classifier_slot_mask_t + num_slots = slot_mask.shape[0] - 1 + slot_int = slot.round().long() + # Out-of-range values land on the sentinel row 0. + safe_idx = slot_int.clamp(min=0, max=num_slots) + is_classifier = (slot_int == safe_idx) & slot_mask[safe_idx] + if not bool(is_classifier.any()): return logits - # Consume it so a same-step prompt_logprobs pass re-applies nothing. - self._pending_verdict = None - rows, verdict, slots, num_reqs = pending - - if logits.shape[0] != num_reqs: - raise RuntimeError( - f"logits rows ({logits.shape[0]}) do not match the forward's " - f"{num_reqs} requests; per-request alignment broke." - ) - # Each marker writes only its own slot's label ids and only that slot's - # real label count (verdict[i, :n]); padded columns are never read. + # Each row writes only its own slot's real label count; padded + # columns are never read. per_slot_label_ids = self.config.classifier_label_token_ids + rows = is_classifier.nonzero(as_tuple=True)[0] logits[rows] = float("-inf") - for i, (r, slot) in enumerate(zip(rows.tolist(), slots.tolist())): - ids = per_slot_label_ids[slot - 1] + for r in rows.tolist(): + row_slot = int(slot_int[r]) + ids = per_slot_label_ids[row_slot - 1] label_ids = torch.tensor(ids, dtype=torch.long, device=logits.device) - logits[r, label_ids] = verdict[i, : len(ids)].to(logits.dtype) + logits[r, label_ids] = verdict[r, : len(ids)].to(logits.dtype) return logits def sample( diff --git a/src/granite_switch/vllm/switch/multi.py b/src/granite_switch/vllm/switch/multi.py index a4606b4e..3e0a1460 100644 --- a/src/granite_switch/vllm/switch/multi.py +++ b/src/granite_switch/vllm/switch/multi.py @@ -105,6 +105,7 @@ from granite_switch.token_exchange import ( apply_token_exchange, build_adapter_kind_lut, + build_classifier_control_luts, build_control_to_substitute_lut, split_adapter_indices, ) @@ -328,15 +329,35 @@ def __init__( self.register_buffer("control_to_substitute_lut", lut, persistent=True) else: self.control_to_substitute_lut = None + marker_lut, slot_lut = build_classifier_control_luts(hf_config) + if marker_lut is not None: + self.register_buffer("classifier_control_lut", marker_lut, persistent=False) + self.register_buffer("classifier_slot_lut", slot_lut, persistent=False) + else: + self.classifier_control_lut = None + self.classifier_slot_lut = None def split_indices( - self, adapter_indices: torch.Tensor + self, + adapter_indices: torch.Tensor, + classifier_probe_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Split ``[total_tokens]`` into (lora_indices, classifier_indices). See :func:`granite_switch.token_exchange.split_adapter_indices`. """ - return split_adapter_indices(self.adapter_kind_lut, adapter_indices) + return split_adapter_indices( + self.adapter_kind_lut, + adapter_indices, + classifier_probe_indices, + ) + + def classifier_indices_from_tokens( + self, input_ids: torch.Tensor + ) -> torch.Tensor | None: + if self.classifier_slot_lut is None: + return None + return self.classifier_slot_lut[input_ids] @property def num_cache_layers(self) -> int: @@ -408,10 +429,16 @@ def forward( # Vectorized control-token matching against ORIGINAL input_ids. matches = input_ids.unsqueeze(1) == adapter_token_ids.unsqueeze(0) # [T, A] is_control_token = matches.any(dim=1) # [total_tokens] + route_control_token = is_control_token + if self.classifier_control_lut is not None: + # A classifier marker is an orthogonal probe. Do not let it replace + # the active adapter stored in the switch's KV-backed coded memory. + is_classifier_marker = self.classifier_control_lut[input_ids] + route_control_token = is_control_token & ~is_classifier_marker # expert_id = argmax + offset (offset selects the layout; see __init__). # 0 for non-control tokens. expert_ids = torch.where( - is_control_token, + route_control_token, matches.long().argmax(dim=1) + self._expert_id_offset, torch.zeros_like(input_ids, dtype=torch.long), ) # [total_tokens] @@ -422,16 +449,16 @@ def forward( _zero = torch.tensor(0.0, dtype=dtype, device=device) _one = torch.tensor(1.0, dtype=dtype, device=device) - # Keys: default masked (-1e9); un-mask (0) at anchor + control tokens - # so the one-hot query attends only to those. Arithmetic (torch.where), - # no boolean-index assignment. + # Keys: default masked (-1e9); un-mask (0) at the anchor and adapter + # selection controls. Classifier markers do not advance the persistent + # adapter address. k_count = torch.full( (total_tokens, self.num_kv_heads, self.counting_head_dim), _NEG_INF, device=device, dtype=dtype, ) - anchor_or_control = is_counting_anchor | is_control_token + anchor_or_control = is_counting_anchor | route_control_token k_count[:, 0, 0] = torch.where(anchor_or_control, _zero, k_count[:, 0, 0]) # Values: v=1 at the position-0 anchor only (the 1/(1+n) numerator). @@ -475,7 +502,7 @@ def forward( # ================================================================== # Look up the codeword for each token's address (compile-safe gather). all_code_vectors = self.codebook[write_addresses].to(dtype) # [T, memory_dim] - write_mask = is_control_token.unsqueeze(-1).to(dtype) # [T, 1] + write_mask = route_control_token.unsqueeze(-1).to(dtype) # [T, 1] # Keys: code(n) * memory_gain at control tokens, zero elsewhere # (arithmetic masking). Value: expert_id at control tokens. @@ -493,7 +520,7 @@ def forward( device=device, dtype=dtype, ) - v_memory[:, 0, 0] = expert_ids.to(dtype) * is_control_token.to(dtype) + v_memory[:, 0, 0] = expert_ids.to(dtype) * route_control_token.to(dtype) # Query: code(n) for every token (reads back the value at address n). q_memory = torch.zeros( diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py index c3999d36..bd619022 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_classifier_compose.py @@ -229,7 +229,7 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): - """[lora, classifier]: LoRA slot 0 populated, classifier slot 1 loaded.""" + """[lora, classifier]: retain LoRA at the classifier marker and score it.""" tokenizer = _tokenizer(MIXED_BASE) hidden = _hidden_size(MIXED_BASE) lora_dir = _find_rag_lora_dir() @@ -292,14 +292,18 @@ def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): "classifier slot 1 LoRA row is not zero — slot numbering leaked" ) - # A detect forward on the classifier control token emits the winning label. - input_ids = _classifier_prompt_ids(tokenizer, clf_ctrl).to( - next(model.parameters()).device - ) + # Select LoRA slot 1, then probe at the classifier's final marker. The + # classifier marker must not replace the active LoRA in the decoder stream. + prompt = _classifier_prompt_ids(tokenizer, clf_ctrl) + input_ids = torch.cat( + [prompt[:, :-1], torch.tensor([[lora_ctrl]]), prompt[:, -1:]], dim=1 + ).to(next(model.parameters()).device) with torch.no_grad(): out = model(input_ids=input_ids, use_cache=False) emitted = int(out.logits[0, -1].argmax().item()) assert emitted == label_ids[winner] + assert model.model._last_lora_indices[0, -1].item() == 1 + assert model.model._last_classifier_indices[0, -1].item() == 2 def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): diff --git a/tests/hf/test_classifier.py b/tests/hf/test_classifier.py index ce7a9b30..e94baacf 100644 --- a/tests/hf/test_classifier.py +++ b/tests/hf/test_classifier.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Classifier-substitute path tests for GraniteSwitchForCausalLM (HF backend). -A classifier slot is an alternative to a LoRA adapter: it shares the adapter -index/control-token machinery, the switch splits classifier positions out of -the LoRA stream, and the classifier head reads the final hidden state to emit -ONE verdict per request. The verdict is read at the control token, which must be -the request's last token. +A classifier slot is selected independently from the active adapter. A LoRA +selected earlier in the prompt remains active at the classifier marker, while +the classifier head reads its slot's read layer (the final hidden state by +default) to emit ONE verdict per request. The verdict marker must be the +request's last token. The verdict exits as a GENERATED LABEL WORD, matching the vLLM backend: the LM head rewrites the marker's logit row so the vocab is -inf except at @@ -25,7 +25,11 @@ def _set_adapter_token_ids(model, token_ids): model.model.adapter_token_ids.data = torch.tensor(token_ids, dtype=torch.long) -def _classifier_config(adapter_kinds, num_labels=2, label_token_ids=(100, 200)): +def _classifier_config( + adapter_kinds, + num_labels=2, + label_token_ids=(100, 200), +): # Full-width label token ids (length num_adapters, None on LoRA slots). Every # classifier slot gets the same ids here; mixed counts are covered by the # compose tests. @@ -34,7 +38,7 @@ def _classifier_config(adapter_kinds, num_labels=2, label_token_ids=(100, 200)): vocab_size=300, hidden_size=64, intermediate_size=128, - num_hidden_layers=3, # 1 switch + 2 decoder + num_hidden_layers=4, # 2 switch cache layers + 2 decoder layers num_attention_heads=4, num_key_value_heads=4, num_adapters=2, @@ -93,6 +97,27 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): assert int(last_row.argmax()) in label_ids +def test_classifier_marker_preserves_active_lora(): + config = _classifier_config(["lora", "classifier"]) + model = GraniteSwitchForCausalLM(config).eval() + _set_adapter_token_ids(model, config.adapter_token_ids) + + # Select LoRA slot 1, then place the classifier marker last. The marker + # probes slot 2 while LoRA slot 1 remains active through the final token. + input_ids = torch.tensor([[10, 250, 20, 30, 251]]) + with torch.no_grad(): + output = model(input_ids=input_ids) + + assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1, 1]] + assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 0, 2]] + label_ids = config.classifier_label_token_ids[1] + last_logits = output.logits[0, -1] + assert torch.isfinite(last_logits[torch.tensor(label_ids)]).all() + assert torch.isneginf( + last_logits[[i for i in range(config.vocab_size) if i not in label_ids]] + ).all() + + def test_non_classifier_request_logits_untouched(): # No classifier token (251) present -> last token is base (index 0), the # head's last-token verdict row is zero, and NO row is blanked: a @@ -280,3 +305,88 @@ def test_right_padded_row_rewrites_the_marker_column(): with pytest.raises(RuntimeError, match="logits_to_keep"): with torch.no_grad(): model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1) + + +def _randomize_classifier_head(model): + """The head bank starts at zero; give it weights so the read layer matters.""" + generator = torch.Generator().manual_seed(0) + head = model.model.classifier_head + with torch.no_grad(): + head.weight.copy_(torch.randn(head.weight.shape, generator=generator)) + head.bias.copy_(torch.randn(head.bias.shape, generator=generator)) + + +def test_intermediate_layer_classifier_scores_that_layers_output(): + """An intermediate read scores the marker row at the configured layer.""" + config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + config.classifier_read_layers = [None, 0] # slot 2: an intermediate layer + model = GraniteSwitchForCausalLM(config).eval() + _set_adapter_token_ids(model, config.adapter_token_ids) + _randomize_classifier_head(model) + with torch.no_grad(): + out = model( + input_ids=torch.tensor([[10, 250, 30, 40, 50, 60, 70, 251]]), + output_hidden_states=True, + ) + + # The classifier marker selects its head independently; the active LoRA + # remains routed through the layer whose residual representation is read. + assert model.model._last_lora_indices[0, -1].item() == 1 + assert model.model._last_classifier_indices[0, -1].item() == 2 + rows, pos, verdict, slots, _ = model.model._last_classifier_verdict + expected = model.model.classifier_head(out.hidden_states[1][rows, pos], slots) + final = model.model.classifier_head(out.hidden_states[-1][rows, pos], slots) + torch.testing.assert_close(verdict, expected) + assert not torch.allclose(verdict, final) + torch.testing.assert_close(out.logits[0, -1, [100, 200]], verdict[0]) + + +def test_final_decoder_layer_read_uses_pre_norm_output(): + """The last decoder layer's residual read precedes the final model norm.""" + config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + config.classifier_read_layers = [None, 1] + model = GraniteSwitchForCausalLM(config).eval() + _set_adapter_token_ids(model, config.adapter_token_ids) + _randomize_classifier_head(model) + captured = [] + hook = model.model.layers[1].register_forward_hook( + lambda _module, _args, output: captured.append(output[0].detach()) + ) + try: + with torch.no_grad(): + out = model( + input_ids=torch.tensor([[10, 20, 30, 40, 50, 60, 70, 251]]), + output_hidden_states=True, + ) + finally: + hook.remove() + + rows, pos, verdict, slots, _ = model.model._last_classifier_verdict + expected = model.model.classifier_head(captured[0][rows, pos], slots) + final_norm = model.model.classifier_head(out.hidden_states[-1][rows, pos], slots) + torch.testing.assert_close(verdict, expected) + assert not torch.allclose(verdict, final_norm) + + +def test_batch_slots_score_at_their_own_read_layers(): + config = _classifier_config(["classifier", "classifier"]) + config.classifier_read_layers = [None, 0] + model = GraniteSwitchForCausalLM(config).eval() + _set_adapter_token_ids(model, config.adapter_token_ids) + _randomize_classifier_head(model) + with torch.no_grad(): + out = model( + input_ids=torch.tensor([[10, 20, 250], [10, 20, 251]]), + output_hidden_states=True, + ) + + rows, pos, verdict, slots, _ = model.model._last_classifier_verdict + assert slots.tolist() == [1, 2] + final_score = model.model.classifier_head( + out.hidden_states[-1][rows[:1], pos[:1]], slots[:1] + ) + intermediate_score = model.model.classifier_head( + out.hidden_states[1][rows[1:], pos[1:]], slots[1:] + ) + torch.testing.assert_close(verdict[:1], final_score) + torch.testing.assert_close(verdict[1:], intermediate_score) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index bfc8b780..2ad4393d 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -258,3 +258,72 @@ def test_lora_targets_include_shared_mlp_groups_when_present(self): "shared_input_linear", "shared_output_linear", } + + +# ════════════════════════════════════════════════════════════════════ +# N. classifier_read_layers +# ════════════════════════════════════════════════════════════════════ + + +def _classifier_kwargs(num_adapters=2, num_decoder_layers=4, **overrides): + """A LoRA slot(s) then one classifier slot, over ``num_decoder_layers``.""" + from granite_switch.config import SWITCH_CACHE_LAYERS + + base = _valid_kwargs( + num_adapters=num_adapters, + num_hidden_layers=num_decoder_layers + SWITCH_CACHE_LAYERS, + adapter_kinds=["lora"] * (num_adapters - 1) + ["classifier"], + classifier_label_token_ids=[None] * (num_adapters - 1) + [[100, 200]], + ) + base.update(overrides) + return base + + +class TestClassifierReadLayers: + def test_omitted_reads_the_final_layer(self): + assert ( + GraniteSwitchConfig(**_classifier_kwargs()).classifier_read_layers is None + ) + + @pytest.mark.parametrize("layers", [[None, None], [None, 0], [None, 3]]) + def test_valid_layers_are_accepted(self, layers): + cfg = GraniteSwitchConfig(**_classifier_kwargs(classifier_read_layers=layers)) + assert cfg.classifier_read_layers == layers + + def test_wrong_length_raises(self): + with pytest.raises(ValueError, match="classifier_read_layers length"): + GraniteSwitchConfig( + **_classifier_kwargs(classifier_read_layers=[None, 2, 3]) + ) + + def test_non_classifier_slot_with_a_layer_raises(self): + with pytest.raises(ValueError, match="non-classifier slot 0"): + GraniteSwitchConfig(**_classifier_kwargs(classifier_read_layers=[1, 2])) + + @pytest.mark.parametrize("layer", [-1, 4, 99, True, 1.0]) + def test_out_of_range_or_non_int_layer_raises(self, layer): + # 4 == num_decoder_layers: the final state is requested by omitting it. + with pytest.raises(ValueError, match=r"in 0\.\.3"): + GraniteSwitchConfig( + **_classifier_kwargs(classifier_read_layers=[None, layer]) + ) + + def test_dual_stream_with_intermediate_layer_raises(self): + with pytest.raises(ValueError, match="dual_stream=True"): + GraniteSwitchConfig( + **_classifier_kwargs( + classifier_read_layers=[None, 2], + dual_stream=True, + cross_stream_rank=8, + ) + ) + + def test_classifier_slot_id_above_256_raises(self): + with pytest.raises(ValueError, match="<= 256"): + GraniteSwitchConfig(**_classifier_kwargs(num_adapters=257)) + + def test_classifier_slot_id_256_is_accepted(self): + GraniteSwitchConfig(**_classifier_kwargs(num_adapters=256)) + + def test_dual_stream_with_final_layer_is_accepted(self): + GraniteSwitchConfig(**_classifier_kwargs(dual_stream=True, cross_stream_rank=8)) diff --git a/tests/unit/test_token_exchange.py b/tests/unit/test_token_exchange.py index f008046c..bde219f9 100644 --- a/tests/unit/test_token_exchange.py +++ b/tests/unit/test_token_exchange.py @@ -6,8 +6,14 @@ """ import pytest +import torch from granite_switch.config import GraniteSwitchConfig +from granite_switch.token_exchange import ( + build_adapter_kind_lut, + build_classifier_control_luts, + split_adapter_indices, +) def _base(num_adapters=2, **overrides): @@ -60,3 +66,40 @@ class TestProjectionHeadDim: def test_inferred_from_hidden_size(self): cfg = GraniteSwitchConfig(**_base()) assert cfg.projection_head_dim == cfg.hidden_size // cfg.num_attention_heads + + +class TestClassifierLoRACoactivation: + def _config(self): + return GraniteSwitchConfig( + **_base( + adapter_token_ids=[499, 500, 501], # leading base-reset token + adapter_substitute_token_ids=[1, 1, 1], + adapter_kinds=["lora", "classifier"], + classifier_label_token_ids=[None, [100, 101]], + ) + ) + + def test_classifier_probe_lut_skips_leading_base_reset_token(self): + cfg = self._config() + marker_lut, slot_lut = build_classifier_control_luts(cfg) + + assert not marker_lut[499] + assert not marker_lut[500] + assert marker_lut[501] + assert slot_lut[499] == slot_lut[500] == 0 + assert slot_lut[501] == 2 + assert cfg.classifier_control_token_ids == [501] + + def test_classifier_probe_preserves_lora_indices_and_separates_probe(self): + cfg = self._config() + adapter_indices = torch.tensor([0, 1, 1, 1]) + probe_indices = torch.tensor([0, 0, 0, 2]) + + lora_indices, classifier_indices = split_adapter_indices( + build_adapter_kind_lut(cfg), + adapter_indices, + probe_indices, + ) + + assert lora_indices.tolist() == [0, 1, 1, 1] + assert classifier_indices.tolist() == [0, 0, 0, 2] diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py index f0e9ac5c..d70860b4 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -96,6 +96,7 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): ckpt = os.path.join(outdir, "composed") model.save_pretrained(ckpt) tok.save_pretrained(ckpt) + del model words = ( "machine translation quietly reshaped how distant villages argued " @@ -125,6 +126,404 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): } +def _reference_hidden(model, prompt_ids, read_layer=None): + """Read the marker's raw decoder output, or final post-norm output. + + HF replaces hidden_states[-1] with the normalized final output, so a + decoder hook is necessary for the last raw layer. Save only the marker + row and always remove the hook, including when forward raises. + """ + import torch + + captured = [] + handle = None + if read_layer is not None: + + def capture(_module, _args, output): + hidden = output[0] if isinstance(output, tuple) else output + captured.append(hidden[0, -1].detach().clone()) + + handle = model.model.layers[read_layer].register_forward_hook(capture) + try: + with torch.no_grad(): + output = model.model(input_ids=torch.tensor([prompt_ids]), use_cache=False) + return ( + captured[0] if read_layer is not None else output.last_hidden_state[0, -1] + ) + finally: + if handle is not None: + handle.remove() + + +def _pick_labels(tok, num_labels): + """The first ``num_labels`` single-token words from ``_LABEL_CANDIDATES``.""" + labels, ids = [], [] + for c in _LABEL_CANDIDATES: + e = tok.encode(c, add_special_tokens=False) + if len(e) == 1 and e[0] not in ids: + labels.append(c) + ids.append(e[0]) + if len(labels) == num_labels: + return labels, ids + raise RuntimeError(f"only {len(labels)} single-token labels, need {num_labels}") + + +def _fixed_head(seed, num_labels, hidden_size): + """Fixed seed-``seed`` weight (std 0.5) and zero bias. Do not search seeds.""" + import torch + + g = torch.Generator().manual_seed(seed) + weight = torch.randn((num_labels, hidden_size), generator=g, dtype=torch.float32) + return weight * 0.5, torch.zeros(num_labels, dtype=torch.float32) + + +def _write_head(slot_dir, weight, bias, labels): + from safetensors.torch import save_file + + from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE + + os.makedirs(slot_dir, exist_ok=True) + save_file( + {"weight": weight.contiguous(), "bias": bias.contiguous()}, + os.path.join(slot_dir, CLASSIFIER_HEAD_FILE), + ) + json.dump( + {"kind": "classifier", "labels": labels}, + open(os.path.join(slot_dir, "adapter_config.json"), "w"), + ) + open(os.path.join(slot_dir, "io.yaml"), "w").write( + f"name: {os.path.basename(slot_dir)}\nmodel: ~\n" + ) + + +def _prompt(tok, control_id): + return [*tok.encode("Is this text safe?", add_special_tokens=False), control_id] + + +def compose_at_layer(base, outdir, read_layer, control_offset=5, num_labels=2): + """Compose one classifier slot reading ``read_layer``, plus its fp32 HF reference.""" + import torch + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + from granite_switch.composer.compose_utils import GraniteSwitchComposer + + tok = AutoTokenizer.from_pretrained(base) + labels, ids = _pick_labels(tok, num_labels) + control_id = tok.vocab_size - control_offset + prompt_ids = _prompt(tok, control_id) + + plain = AutoModelForCausalLM.from_pretrained(base, dtype=torch.float32).eval() + reference_hidden = _reference_hidden(plain, prompt_ids, read_layer) + del plain + + weight, bias = _fixed_head( + 0, len(ids), AutoConfig.from_pretrained(base).hidden_size + ) + slot = os.path.join(outdir, "detect") + _write_head(slot, weight, bias, labels) + + model = GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=base, + adapter_paths=[slot], + adapter_token_ids=[control_id], + adapter_substitute_token_ids=[control_id], + adapter_names=["detect"], + adapter_kinds=["classifier"], + classifier_label_token_ids=[ids], + classifier_read_layers=[read_layer], + ) + ckpt = os.path.join(outdir, "composed") + model.save_pretrained(ckpt) + tok.save_pretrained(ckpt) + del model + + scores = weight.double() @ reference_hidden.double() + bias.double() + return { + "ckpt": ckpt, + "control_id": control_id, + "label_token_ids": ids, + "prompt_ids": prompt_ids, + "read_layer": read_layer, + "ground_truth_label_id": ids[int(scores.argmax())], + "ground_truth_scores": scores.tolist(), + } + + +def compose_two_slot(base, outdir, read_layer, control_offset=5, num_labels=2): + """Compose a final-layer and a ``read_layer`` slot, plus fp32 HF references.""" + import torch + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + from granite_switch.composer.compose_utils import GraniteSwitchComposer + + tok = AutoTokenizer.from_pretrained(base) + labels, ids = _pick_labels(tok, num_labels) + hidden_size = AutoConfig.from_pretrained(base).hidden_size + control_ids = [tok.vocab_size - control_offset, tok.vocab_size - control_offset - 1] + heads = [_fixed_head(seed, len(ids), hidden_size) for seed in (0, 1)] + slot_dirs = [ + os.path.join(outdir, "detect_final"), + os.path.join(outdir, "detect_mid"), + ] + for d, (w, b) in zip(slot_dirs, heads): + _write_head(d, w, b, labels) + + model = GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=base, + adapter_paths=slot_dirs, + adapter_token_ids=control_ids, + adapter_substitute_token_ids=control_ids, + adapter_names=["detect_final", "detect_mid"], + adapter_kinds=["classifier", "classifier"], + classifier_label_token_ids=[ids, ids], + classifier_read_layers=[None, read_layer], + ) + ckpt = os.path.join(outdir, "composed_two_slot") + model.save_pretrained(ckpt) + tok.save_pretrained(ckpt) + del model + + plain = AutoModelForCausalLM.from_pretrained(base, dtype=torch.float32).eval() + result = {"ckpt": ckpt, "label_token_ids": ids, "read_layer": read_layer} + for name, control_id, (w, b), layer in zip( + ("slot1", "slot2"), control_ids, heads, (None, read_layer) + ): + prompt_ids = _prompt(tok, control_id) + hidden = _reference_hidden(plain, prompt_ids, layer) + scores = w.double() @ hidden.double() + b.double() + result[f"{name}_prompt_ids"] = prompt_ids + result[f"{name}_reference_hidden"] = hidden.tolist() + result[f"{name}_ground_truth_label_id"] = ids[int(scores.argmax())] + result[f"{name}_ground_truth_scores"] = scores.tolist() + return result + + +def _record_classifier_rows(llm): + """Copy each sampled classifier row as ``compute_logits`` sees it. + + Installed after engine startup, so profiling and capture runs are not + recorded. ``compute_logits`` runs outside the captured forward, and the + wrapper only copies tensors after the original returns. Requires the + in-process engine so the returned list is filled in this process. + """ + from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM + + records = [] + + def install(model): + target = next( + m for m in model.modules() if isinstance(m, GraniteSwitchForCausalLM) + ) + original = target.compute_logits + config = target.config + head = target.model.classifier_head + + def compute_logits(hidden_states): + logits = original(hidden_states) + slots = hidden_states[:, -1].round().long() + for row in (slots > 0).nonzero(as_tuple=True)[0].tolist(): + slot = int(slots[row]) + labels = config.classifier_label_token_ids[slot - 1] + n = len(labels) + h = config.hidden_size + record = { + "slot": slot, + "packed_scores": hidden_states[row, h : h + n].float().cpu(), + "rewritten_label_logits": logits[row, labels].float().cpu(), + } + # The packed hidden-state prefix is the final post-norm + # activation: the actual head input only for final-layer slots. + if config.classifier_read_layers[slot - 1] is None: + record["served_hidden"] = hidden_states[row, :h].float().cpu() + record["loaded_weight"] = head.weight[slot - 1, :n].float().cpu() + record["loaded_bias"] = head.bias[slot - 1, :n].float().cpu() + records.append(record) + return logits + + target.compute_logits = compute_logits + + llm.apply_model(install) + return records + + +def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, eager): + """Serve both classifier slots' prompts, one request each, eager or FULL. + + Under FULL both requests share one captured batch shape, exercising + replay with two different slots/read layers. Reports per-label logprobs + and the recorded classifier row (packed scores, rewritten label logits, + and for the final-layer slot its actual activation and loaded head). + """ + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + from vllm import LLM, SamplingParams + from vllm.inputs import TokensPrompt + + kw = dict( + model=ckpt, + enforce_eager=eager, + gpu_memory_utilization=0.55, + max_model_len=4096, + enable_prefix_caching=False, + dtype="bfloat16", + ) + if not eager: + kw["compilation_config"] = {"cudagraph_mode": "FULL"} + llm = LLM(**kw) + records = _record_classifier_rows(llm) + sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=len(label_token_ids)) + + results = [] + for expected_slot, prompt_ids in ((1, prompt_ids_1), (2, prompt_ids_2)): + records.clear() + out = llm.generate([TokensPrompt(prompt_token_ids=prompt_ids)], sp)[0].outputs[ + 0 + ] + recorded_slots = [r["slot"] for r in records] + if recorded_slots != [expected_slot]: + raise RuntimeError( + f"expected one recorded slot-{expected_slot} row, got {recorded_slots}" + ) + lp_dict = out.logprobs[0] + label_logprobs = [ + lp_dict[tid].logprob if tid in lp_dict else None for tid in label_token_ids + ] + results.append( + { + "token_id": int(out.token_ids[0]), + "text": out.text, + "label_logprobs": label_logprobs, + **{k: v.tolist() for k, v in records[0].items() if k != "slot"}, + } + ) + + json.dump({"slot1": results[0], "slot2": results[1]}, open(out_path, "w")) + + +def run_at_layer(ckpt, prompt_ids, label_token_ids, out_path, full_cudagraph=False): + """Serve one prompt; report the emitted token and every label's logprob.""" + from vllm import LLM, SamplingParams + from vllm.inputs import TokensPrompt + + kw = dict( + model=ckpt, + gpu_memory_utilization=0.55, + max_model_len=4096, + enable_prefix_caching=False, + dtype="bfloat16", + ) + if full_cudagraph: + kw["enforce_eager"] = False + kw["compilation_config"] = {"cudagraph_mode": "FULL"} + else: + kw["enforce_eager"] = True + llm = LLM(**kw) + sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=len(label_token_ids)) + out = llm.generate([TokensPrompt(prompt_token_ids=prompt_ids)], sp)[0].outputs[0] + lp_dict = out.logprobs[0] + label_logprobs = [ + lp_dict[tid].logprob if tid in lp_dict else None for tid in label_token_ids + ] + json.dump( + { + "token_id": int(out.token_ids[0]), + "text": out.text, + "label_logprobs": label_logprobs, + }, + open(out_path, "w"), + ) + + +def run_full_cudagraph_replay( + ckpt, control_id, label_token_ids, out_path, *, eager=False +): + """Run identical fixed-length prompts in eager or FULL mode. + + Includes six changed inputs, a plain request, and a mixed batch. Markers + remain last. Equal prompt lengths permit graph reuse; this test does not + instrument the dispatcher to count actual graph replays. + """ + import math + import random + + from vllm import LLM, SamplingParams + from vllm.inputs import TokensPrompt + + label_ids = set(label_token_ids) + kw = dict( + model=ckpt, + enforce_eager=eager, + gpu_memory_utilization=0.55, + max_model_len=4096, + enable_prefix_caching=False, + dtype="bfloat16", + ) + if not eager: + kw["compilation_config"] = {"cudagraph_mode": "FULL"} + llm = LLM(**kw) + sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=len(label_ids) + 1) + + def summarize(output): + return { + "token_id": int(output.token_ids[0]), + "text": output.text, + "is_label": int(output.token_ids[0]) in label_ids, + "has_non_label_support": any( + tid not in label_ids and math.isfinite(lp.logprob) + for tid, lp in output.logprobs[0].items() + ), + } + + rng = random.Random(0) + base_len = 40 + results = [] + for _ in range(6): + body = [rng.randint(1000, 2000) for _ in range(base_len)] + prompt = TokensPrompt(prompt_token_ids=[*body, control_id]) + o = llm.generate([prompt], sp)[0].outputs[0] + results.append(summarize(o)) + + # Same shape (base_len + 1 tokens), no classifier marker at all -- the + # verdict from the PRECEDING call must not leak into this one. This is a + # SEQUENTIAL check (separate generate() calls), not a same-batch mix. + plain_body = [rng.randint(1000, 2000) for _ in range(base_len + 1)] + plain_prompt = TokensPrompt(prompt_token_ids=plain_body) + o = llm.generate([plain_prompt], sp)[0].outputs[0] + plain_result = summarize(o) + + # A GENUINE mixed batch: 3 requests in ONE generate() call -- classifier, + # plain, classifier -- so vLLM schedules and samples all three rows in the + # same forward pass. Each must get its own correct result: the plain + # row's logits must not pick up either classifier row's verdict, and the + # two classifier rows (same slot here, since this checkpoint has one) + # must each still fire independently. + mix_cls_body_1 = [rng.randint(1000, 2000) for _ in range(base_len)] + mix_plain_body = [rng.randint(1000, 2000) for _ in range(base_len + 1)] + mix_cls_body_2 = [rng.randint(1000, 2000) for _ in range(base_len)] + mix_prompts = [ + TokensPrompt(prompt_token_ids=[*mix_cls_body_1, control_id]), + TokensPrompt(prompt_token_ids=mix_plain_body), + TokensPrompt(prompt_token_ids=[*mix_cls_body_2, control_id]), + ] + mix_outputs = llm.generate(mix_prompts, sp) + mixed_results = [summarize(o.outputs[0]) for o in mix_outputs] + + json.dump( + { + "results": results, + "plain_result": plain_result, + "mixed_results": mixed_results, + }, + open(out_path, "w"), + ) + + +def run_eager_reference(ckpt, control_id, out_path): + # Reuse the exact prompt sequence, including the plain and mixed requests. + run_full_cudagraph_replay( + ckpt, control_id, json.loads(os.environ["CLS_LABEL_IDS"]), out_path, eager=True + ) + + def run_budget(ckpt, prompt_ids, budget, out_path): from vllm import LLM, SamplingParams from vllm.inputs import TokensPrompt @@ -168,6 +567,91 @@ def main(): json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) sys.exit(1) return + if mode == "run-full-cudagraph": + out_path = sys.argv[2] + control_id = int(os.environ["CLS_CONTROL_ID"]) + label_token_ids = json.loads(os.environ["CLS_LABEL_IDS"]) + try: + run_full_cudagraph_replay( + os.environ["CLS_CKPT"], control_id, label_token_ids, out_path + ) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return + if mode == "run-eager-reference": + out_path = sys.argv[2] + control_id = int(os.environ["CLS_CONTROL_ID"]) + try: + run_eager_reference(os.environ["CLS_CKPT"], control_id, out_path) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return + if mode == "compose-at-layer": + read_layer = int(sys.argv[2]) + try: + meta = compose_at_layer( + os.environ["CLS_BASE"], os.environ["CLS_OUT"], read_layer + ) + print(json.dumps({"result": meta})) + except Exception as e: + print(json.dumps({"error": f"{e}\n{traceback.format_exc()}"})) + return + if mode == "run-at-layer": + out_path = sys.argv[2] + prompt_ids = json.load(open(os.environ["CLS_PIDS"])) + label_token_ids = json.loads(os.environ["CLS_LABEL_IDS"]) + try: + run_at_layer(os.environ["CLS_CKPT"], prompt_ids, label_token_ids, out_path) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return + if mode == "run-at-layer-full-cudagraph": + out_path = sys.argv[2] + prompt_ids = json.load(open(os.environ["CLS_PIDS"])) + label_token_ids = json.loads(os.environ["CLS_LABEL_IDS"]) + try: + run_at_layer( + os.environ["CLS_CKPT"], + prompt_ids, + label_token_ids, + out_path, + full_cudagraph=True, + ) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return + if mode == "compose-two-slot": + read_layer = int(sys.argv[2]) + try: + meta = compose_two_slot( + os.environ["CLS_BASE"], os.environ["CLS_OUT"], read_layer + ) + print(json.dumps({"result": meta})) + except Exception as e: + print(json.dumps({"error": f"{e}\n{traceback.format_exc()}"})) + return + if mode in ("run-two-slot-eager", "run-two-slot-full-cudagraph"): + out_path = sys.argv[2] + prompt_ids_1 = json.loads(os.environ["CLS_PIDS_1"]) + prompt_ids_2 = json.loads(os.environ["CLS_PIDS_2"]) + label_token_ids = json.loads(os.environ["CLS_LABEL_IDS"]) + try: + run_two_slot( + os.environ["CLS_CKPT"], + prompt_ids_1, + prompt_ids_2, + label_token_ids, + out_path, + eager=mode == "run-two-slot-eager", + ) + except Exception as e: + json.dump({"fatal": f"{e}\n{traceback.format_exc()}"}, open(out_path, "w")) + sys.exit(1) + return raise SystemExit(f"unknown mode {mode!r}") diff --git a/tests/vllm/_model_forward_tests.py b/tests/vllm/_model_forward_tests.py index 73cba6ff..1ec442dc 100644 --- a/tests/vllm/_model_forward_tests.py +++ b/tests/vllm/_model_forward_tests.py @@ -59,7 +59,10 @@ def _try_import_vllm(): # ── Helpers ────────────────────────────────────────────────────────── -def _tiny_vllm_config(): +def _tiny_vllm_config( + adapter_kinds=None, + classifier_label_token_ids=None, +): """Minimal GraniteSwitchConfig for single-GPU vLLM tests.""" return GraniteSwitchConfig( vocab_size=300, @@ -76,6 +79,8 @@ def _tiny_vllm_config(): # finalize_weights asserts this for applicable adapters. max_lora_rank=16, adapter_ranks=[16, 16], + adapter_kinds=adapter_kinds, + classifier_label_token_ids=classifier_label_token_ids, switch_head_dim=32, max_position_embeddings=512, attention_multiplier=1.0, @@ -184,7 +189,7 @@ def setup_model(self): _ensure_distributed() self.device = torch.device("cuda") - self.config = _tiny_vllm_config() + self.config = self._make_config() old_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.bfloat16) @@ -229,6 +234,9 @@ def setup_model(self): for name in list(self._layer_names): sfc.pop(name, None) + def _make_config(self): + return _tiny_vllm_config() + def _setup_kv_caches(self): self._kv_caches = [] self._attention_map = {} @@ -466,6 +474,42 @@ def test_different_adapters_produce_different_logits(self): ) +class TestClassifierLoRAProbe(_VLLMModelTestBase): + def _make_config(self): + return _tiny_vllm_config( + adapter_kinds=["lora", "classifier"], + classifier_label_token_ids=[None, [100, 101]], + ) + + def test_marker_scores_with_active_lora_and_emits_label(self, monkeypatch): + self.model.eval() + seen = {} + switch = self.model.model.switch + original_split = switch.split_indices + + def record_split(adapter_indices, classifier_probe_indices=None): + lora_indices, classifier_indices = original_split( + adapter_indices, classifier_probe_indices + ) + seen["lora"] = lora_indices.detach().cpu() + seen["classifier"] = classifier_indices.detach().cpu() + return lora_indices, classifier_indices + + monkeypatch.setattr(switch, "split_indices", record_split) + + with torch.no_grad(): + logits = self._run_forward_and_logits([10, 250, 20, 30, 251]) + + assert seen["lora"].tolist() == [0, 1, 1, 1, 1] + assert seen["classifier"].tolist() == [0, 0, 0, 0, 2] + final_logits = logits[-1] + assert torch.isfinite(final_logits[[100, 101]]).all() + non_label_ids = [ + i for i in range(self.config.vocab_size) if i not in (100, 101) + ] + assert torch.isneginf(final_logits[non_label_ids]).all() + + # ════════════════════════════════════════════════════════════════════ # 5. KV visibility tests # ════════════════════════════════════════════════════════════════════ diff --git a/tests/vllm/test_classifier_capture_safety.py b/tests/vllm/test_classifier_capture_safety.py new file mode 100644 index 00000000..c0285dad --- /dev/null +++ b/tests/vllm/test_classifier_capture_safety.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU method tests for token-aligned, score-at-layer classification. + +GPU capture/replay and packed-output transport are tested in +``test_classifier_chunked_prefill.py``. Request-boundary validation is deferred. +""" + +import pytest +import torch + +pytest.importorskip("vllm") + +from granite_switch.vllm.granite_switch_model import ( + GraniteSwitchForCausalLM, + GraniteSwitchModel, +) + + +class _StubClassifierHead: + """Score = slot id * hidden[:, 0], so a score shows which slot and source it read.""" + + def __call__(self, x, classifier_indices): + return (classifier_indices.float() * x[:, 0]).unsqueeze(-1).expand(-1, 2) + + +class _Model(torch.nn.Module): + """Stands in for ``GraniteSwitchModel`` with the attributes its __init__ builds. + + ``read_layers``/``kinds`` are per 1-based slot; None reads the final layer. + """ + + _score_group = GraniteSwitchModel._score_group + + def __init__(self, slot_of_marker, read_layers=None, kinds=None): + super().__init__() + num_slots = max(slot_of_marker.values()) + read_layers = read_layers or [None] * num_slots + is_classifier = [k == "classifier" for k in kinds or ["classifier"] * num_slots] + self.config = type("Cfg", (), {"max_classifier_labels": 2})() + self.classifier_head = _StubClassifierHead() + self.classifier_control_token_ids_t = torch.tensor(sorted(slot_of_marker)) + self.classifier_slot_mask_t = torch.tensor([False, *is_classifier]) + slots_by_layer = {} + for slot, (is_clf, layer) in enumerate( + zip(is_classifier, read_layers), start=1 + ): + if is_clf: + slots_by_layer.setdefault(layer, []).append(slot) + self._classifier_final_layer_group = None + self._classifier_layer_groups = {} + for layer, slot_ids in slots_by_layer.items(): + name = f"_classifier_slots_{'final' if layer is None else layer}" + self.register_buffer(name, torch.tensor(slot_ids)) + if layer is None: + self._classifier_final_layer_group = name + else: + self._classifier_layer_groups[layer] = name + + +def _run(model, slot_of_marker, tokens, input_ids_override=...): + """Call ``_compute_classifier_outputs`` on a flat token list; hidden[t] = t + 1.""" + indices, current = [], 0 + for t in tokens: + current = slot_of_marker.get(t, current) + indices.append(current) + input_ids = ( + torch.tensor(tokens) if input_ids_override is ... else input_ids_override + ) + hidden = torch.arange(1.0, len(tokens) + 1).unsqueeze(1).expand(-1, 4) + return GraniteSwitchModel._compute_classifier_outputs( + model, input_ids, hidden, torch.tensor(indices), torch.zeros(len(tokens), 2) + ) + + +def _score(model, layer, hidden_value, indices, scores): + hidden = torch.full((len(indices), 4), hidden_value) + return GraniteSwitchModel._score_group( + model, + model._classifier_layer_groups[layer], + hidden, + torch.tensor(indices), + scores, + ) + + +def test_every_marker_token_gets_its_slot_and_score(): + model = _Model({90: 2, 91: 3}) + scores, slot = _run(model, {90: 2, 91: 3}, [5, 6, 90, 7, 8, 9, 91]) + assert slot.tolist() == [0, 0, 2, 0, 0, 0, 3] + assert scores[2].tolist() == [6.0, 6.0] # slot 2 * hidden 3 + assert scores[6].tolist() == [21.0, 21.0] # slot 3 * hidden 7 + + +def test_marker_is_identified_by_id_at_any_position(): + model = _Model({90: 2, 91: 3}) + _, slot = _run(model, {90: 2, 91: 3}, [90, 5, 6, 91]) + assert slot.tolist() == [2, 0, 0, 3] + + +def test_repeated_marker_of_the_same_slot_is_reported_each_time(): + model = _Model({90: 2, 91: 3}) + _, slot = _run(model, {90: 2, 91: 3}, [90, 5, 90, 6]) + assert slot.tolist() == [2, 0, 2, 0] + + +def test_no_marker_is_all_zero(): + model = _Model({90: 2, 91: 3}) + scores, slot = _run(model, {90: 2, 91: 3}, [5, 6, 7]) + assert slot.tolist() == [0, 0, 0] + assert not scores.any() + + +def test_no_input_ids_is_all_zero_slot(): + model = _Model({90: 2, 91: 3}) + _, slot = _run(model, {90: 2, 91: 3}, [5, 90, 6], input_ids_override=None) + assert slot.tolist() == [0, 0, 0] + + +def test_final_layer_scoring_leaves_intermediate_slot_scores_untouched(): + model = _Model({90: 2, 91: 3}, read_layers=[None, 1, None]) + pre_scores = torch.tensor([[0.0, 0.0], [99.0, 99.0], [0.0, 0.0]]) + scores, slot = GraniteSwitchModel._compute_classifier_outputs( + model, + torch.tensor([5, 90, 91]), + torch.ones(3, 4), + torch.tensor([0, 2, 3]), + pre_scores, + ) + assert slot.tolist() == [0, 2, 3] + assert scores[1].tolist() == [99.0, 99.0] # slot 2 reads layer 1 + assert scores[2].tolist() == [3.0, 3.0] + + +def test_each_intermediate_layer_scores_only_its_own_slots(): + model = _Model({90: 1, 91: 2}, read_layers=[3, 7]) + scores = _score(model, 3, 2.0, [0, 1, 2], torch.zeros(3, 2)) + assert scores[1].tolist() == [2.0, 2.0] + assert scores[2].tolist() == [0.0, 0.0] + scores = _score(model, 7, 5.0, [0, 1, 2], scores) + assert scores[1].tolist() == [2.0, 2.0] # layer 3's value kept + assert scores[2].tolist() == [10.0, 10.0] + + +def test_slots_sharing_a_layer_form_one_group(): + model = _Model({90: 1, 91: 2}) + assert model._classifier_layer_groups == {} + assert getattr(model, model._classifier_final_layer_group).tolist() == [1, 2] + scores, slot = _run(model, {90: 1, 91: 2}, [90, 5, 91]) + assert slot.tolist() == [1, 0, 2] + assert scores[0].tolist() == [1.0, 1.0] + assert scores[2].tolist() == [6.0, 6.0] + + +def test_lora_slots_form_no_group(): + model = _Model({91: 2}, kinds=["lora", "classifier"]) + assert getattr(model, model._classifier_final_layer_group).tolist() == [2] + assert model._classifier_layer_groups == {} + _, slot = _run(model, {91: 2}, [5, 91, 6]) + assert slot.tolist() == [0, 2, 0] + + +def test_profiling_random_slot_on_a_lora_slot_leaves_logits_untouched(): + """vLLM's profiling sampler run feeds random hidden states; 0.75 rounds to LoRA slot 1.""" + wrapper = GraniteSwitchForCausalLM.__new__(GraniteSwitchForCausalLM) + torch.nn.Module.__init__(wrapper) + wrapper.model = _Model({91: 2}, kinds=["lora", "classifier"]) + wrapper.config = type( + "Cfg", (), {"classifier_label_token_ids": [None, [100, 200]]} + )() + logits = torch.randn(1, 300) + expected = logits.clone() + out = GraniteSwitchForCausalLM._apply_classifier_verdict( + wrapper, logits, torch.randn(1, 2), torch.tensor([0.75]) + ) + assert torch.equal(out, expected) + + +def test_scoring_resolves_replaced_registered_buffers(): + model = _Model({90: 1, 91: 2}, read_layers=[None, 0]) + old_buffers = list(model.buffers()) + # Device moves replace buffers; poison the old storage to catch stale aliases. + model._apply(lambda tensor: tensor.clone()) + for tensor in old_buffers: + tensor.zero_() + scores = _score(model, 0, 3.0, [1, 2], torch.zeros(2, 2)) + scores, slots = GraniteSwitchModel._compute_classifier_outputs( + model, + torch.tensor([90, 91]), + torch.full((2, 4), 5.0), + torch.tensor([1, 2]), + scores, + ) + assert scores.tolist() == [[5.0, 5.0], [6.0, 6.0]] + assert slots.tolist() == [1.0, 2.0] diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_classifier_chunked_prefill.py index 4590a41e..6c7e3519 100644 --- a/tests/vllm/test_classifier_chunked_prefill.py +++ b/tests/vllm/test_classifier_chunked_prefill.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""Chunked-prefill coverage for the vLLM classifier verdict exit. +"""GPU coverage for the vLLM classifier verdict exit: FULL capture/replay, +chunked prefill, intermediate-layer reads, and two slots in one checkpoint. ``enable_chunked_prefill`` is on by default in serving. The prompt is laid out as the chat template renders it, with the classifier control token as the very last @@ -48,11 +49,11 @@ def _env(): return env -def _compose(tmp_path_factory, mode, tag): +def _compose(tmp_path_factory, argv, tag): """Compose a classifier checkpoint in a CPU subprocess, before any CUDA init.""" outdir = str(tmp_path_factory.mktemp(tag)) proc = subprocess.run( - [sys.executable, WORKER, mode], + [sys.executable, WORKER, *argv], env={**_env(), "CLS_BASE": BASE_MODEL, "CLS_OUT": outdir}, capture_output=True, text=True, @@ -61,40 +62,35 @@ def _compose(tmp_path_factory, mode, tag): lines = [l for l in (proc.stdout or "").splitlines() if l.startswith("{")] if proc.returncode != 0 or not lines: pytest.fail( - f"{mode} failed:\nstdout={proc.stdout[-2000:]}\n" + f"{argv[0]} failed:\nstdout={proc.stdout[-2000:]}\n" f"stderr={proc.stderr[-2000:]}", pytrace=False, ) meta = json.loads(lines[-1]) if "error" in meta: - pytest.fail(f"{mode} error: {meta['error'][:2000]}", pytrace=False) + pytest.fail(f"{argv[0]} error: {meta['error'][:2000]}", pytrace=False) result = meta["result"] - result["pids_path"] = os.path.join(outdir, "pids.json") - json.dump(result["prompt_ids"], open(result["pids_path"], "w")) result["outdir"] = outdir return result -@pytest.fixture(scope="module") -def composed(tmp_path_factory): - """Multi-label classifier over a varied prompt.""" - return _compose(tmp_path_factory, "compose-multilabel", "cls_ckpt") +def _write_pids(result): + path = os.path.join(result["outdir"], "pids.json") + json.dump(result["prompt_ids"], open(path, "w")) + return path -def _run_worker(composed, argv, out_name, label, timeout): +def _run_worker(outdir, argv, env, timeout=1800): """Run one worker invocation in its own subprocess; return the parsed result.""" - out_path = os.path.join(composed["outdir"], out_name) + out_path = os.path.join(outdir, "_".join(argv) + ".json") proc = subprocess.run( [sys.executable, WORKER, *argv, out_path], - env={ - **_env(), - "CLS_CKPT": composed["ckpt"], - "CLS_PIDS": composed["pids_path"], - }, + env={**_env(), **env}, capture_output=True, text=True, timeout=timeout, ) + label = " ".join(argv) if not os.path.exists(out_path): pytest.fail( f"{label}: worker wrote no result\n" @@ -107,12 +103,56 @@ def _run_worker(composed, argv, out_name, label, timeout): return res -def _run(composed, budget): - """Serve the prompt at one chunk budget (``None`` = chunked prefill disabled).""" - tag = "none" if budget is None else str(budget) - return _run_worker( - composed, ["run", tag], f"run_{tag}.json", f"budget={budget}", 1200 +@pytest.fixture(scope="module") +def composed(tmp_path_factory): + """Multi-label classifier over a varied prompt.""" + result = _compose(tmp_path_factory, ["compose-multilabel"], "cls_ckpt") + result["pids_path"] = _write_pids(result) + return result + + +def _composed_env(composed, **extra): + return {"CLS_CKPT": composed["ckpt"], "CLS_PIDS": composed["pids_path"], **extra} + + +def test_full_cudagraph_mode_does_not_crash_and_matches_eager(composed): + """Exercise FULL startup and changed inputs, including mixed requests. + + The original path copied marker IDs to CUDA during capture and crashed. + Compare the selected tokens with the same prompts served eagerly. + """ + label_ids = set(composed["label_token_ids"]) + env = _composed_env( + composed, + CLS_CONTROL_ID=str(composed["control_id"]), + CLS_LABEL_IDS=json.dumps(composed["label_token_ids"]), ) + full = _run_worker(composed["outdir"], ["run-full-cudagraph"], env) + eager = _run_worker(composed["outdir"], ["run-eager-reference"], env) + + assert len(full["results"]) == len(eager["results"]) == 6 + for i, (f, e) in enumerate(zip(full["results"], eager["results"])): + assert f["is_label"], ( + f"call {i}: FULL-cudagraph verdict {f['token_id']} ({f['text']!r}) " + f"is not a label token {label_ids}; the classifier exit did not " + f"fire under real CUDA-graph capture/replay" + ) + assert f["token_id"] == e["token_id"], ( + f"call {i}: FULL-cudagraph emitted {f['token_id']} ({f['text']!r}) " + f"!= eager reference {e['token_id']} ({e['text']!r})" + ) + + # Ordinary generation may legitimately choose a label word. Check that + # non-label tokens remain possible, rather than banning those words. + assert full["plain_result"]["has_non_label_support"] + assert full["plain_result"]["token_id"] == eager["plain_result"]["token_id"] + mixed = full["mixed_results"] + assert len(mixed) == len(eager["mixed_results"]) == 3 + for i, (actual, reference) in enumerate(zip(mixed, eager["mixed_results"])): + assert actual["token_id"] == reference["token_id"], f"mixed row {i}" + assert actual["has_non_label_support"] == (i == 1) + if i != 1: + assert actual["is_label"] def test_classifier_token_stable_across_chunk_budgets(composed): @@ -122,7 +162,15 @@ def test_classifier_token_stable_across_chunk_budgets(composed): "the prompt must be longer than the smallest budget so it chunks" ) - results = {b: _run(composed, b) for b in BUDGETS} + results = { + b: _run_worker( + composed["outdir"], + ["run", "none" if b is None else str(b)], + _composed_env(composed), + timeout=1200, + ) + for b in BUDGETS + } ref = results[None] assert ref["token_id"] in label_ids, ( @@ -134,3 +182,242 @@ def test_classifier_token_stable_across_chunk_budgets(composed): f"budget={b} emitted {r['token_id']} ({r['text']!r}) != unchunked " f"{ref['token_id']} ({ref['text']!r})" ) + + +# Score-gap tolerance against the fp32 HF reference; this comparison also +# includes HF/vLLM backbone differences, not only classifier arithmetic. +SCORE_DIFF_ATOL = 0.3 + + +def _check_served_scores( + name, served, ref_scores, ref_label_id, label_token_ids, min_label_margin +): + """Print one served verdict beside its CPU reference; return failure messages. + + Reference values are raw head scores; served values are logprobs, so only + label-to-label differences are comparable. The winning-label check runs + only when the reference margin is at least ``min_label_margin`` (``None`` + always runs it); below that a near-tie makes it inconclusive. + """ + ref_sorted = sorted(ref_scores, reverse=True) + margin = ref_sorted[0] - ref_sorted[1] + label_logprobs = served["label_logprobs"] + print(f"\n[{name}] label token ids={label_token_ids}") + print(f" reference raw scores (fp32): {ref_scores}") + print(f" reference label={ref_label_id}, decision margin={margin:.4f}") + print(f" served label logprobs: {label_logprobs}") + print(f" served selected token={served['token_id']} ({served['text']!r})") + + failures = [] + if any(lp is None for lp in label_logprobs): + failures.append( + f"{name}: not every label token appeared in the top-" + f"{len(label_token_ids)} logprobs ({label_logprobs})" + ) + else: + for j in range(1, len(label_token_ids)): + served_diff = label_logprobs[j] - label_logprobs[0] + ref_diff = ref_scores[j] - ref_scores[0] + ok = served_diff == pytest.approx(ref_diff, abs=SCORE_DIFF_ATOL) + print( + f" label[{j}]-label[0]: served logprob diff={served_diff:.4f}, " + f"reference score diff={ref_diff:.4f} -> {'ok' if ok else 'MISMATCH'}" + ) + if not ok: + failures.append( + f"{name}: served logprob diff label[{j}]-label[0] = " + f"{served_diff:.4f} != reference score diff {ref_diff:.4f} " + f"+/- {SCORE_DIFF_ATOL}" + ) + + if min_label_margin is not None and margin < min_label_margin: + print( + f" winning-label check INCONCLUSIVE: reference margin " + f"{margin:.4f} < {min_label_margin} (near-tie)" + ) + elif served["token_id"] != ref_label_id: + failures.append( + f"{name}: served token {served['token_id']} ({served['text']!r}) " + f"!= reference label {ref_label_id} (margin={margin:.4f})" + ) + return failures + + +# Allowance for fp32 log-softmax subtraction only (fp32 spacing is ~6e-5 at +# |logit| < 512); not a numerical-equivalence tolerance. +LOGPROB_DIFF_ATOL = 1e-3 + + +def _check_transport(name, served, label_token_ids): + """Packed scores -> rewritten label logits -> returned logprobs.""" + packed = served["packed_scores"] + rewritten = served["rewritten_label_logits"] + label_logprobs = served["label_logprobs"] + print(f"\n[{name}] transport") + print(f" packed head scores (raw): {packed}") + print(f" rewritten label logits (raw): {rewritten}") + print(f" served label logprobs: {label_logprobs}") + + failures = [] + if rewritten != packed: + failures.append( + f"{name}: rewritten label logits {rewritten} != packed scores {packed}" + ) + if any(lp is None for lp in label_logprobs): + failures.append( + f"{name}: not every label token appeared in the top-" + f"{len(label_token_ids)} logprobs ({label_logprobs})" + ) + return failures + for j in range(1, len(label_token_ids)): + lp_diff = label_logprobs[j] - label_logprobs[0] + logit_diff = rewritten[j] - rewritten[0] + if lp_diff != pytest.approx(logit_diff, abs=LOGPROB_DIFF_ATOL): + failures.append( + f"{name}: logprob diff label[{j}]-label[0] = {lp_diff} != " + f"rewritten logit diff {logit_diff} +/- {LOGPROB_DIFF_ATOL}" + ) + return failures + + +def _check_final_layer_arithmetic(name, served): + """Packed scores against the loaded head applied to the actual activation.""" + weight = torch.tensor(served["loaded_weight"]) + bias = torch.tensor(served["loaded_bias"]) + hidden = torch.tensor(served["served_hidden"]) + recomputed = weight @ hidden + bias + rounded = recomputed.to(torch.bfloat16).float().tolist() + print(f"\n[{name}] arithmetic on the served activation") + print(f" loaded head in fp32: {recomputed.tolist()}") + print(f" ... rounded to bf16: {rounded}") + print(f" packed head scores (raw): {served['packed_scores']}") + if rounded != served["packed_scores"]: + return [ + f"{name}: packed scores {served['packed_scores']} != loaded head " + f"on served activation, rounded to bf16, {rounded}" + ] + return [] + + +def _print_final_layer_activation(name, served, reference_hidden, reference_scores): + """Record the HF/vLLM final-activation difference; no acceptance threshold.""" + weight = torch.tensor(served["loaded_weight"]) + bias = torch.tensor(served["loaded_bias"]) + hf = torch.tensor(reference_hidden, dtype=torch.float32) + vllm = torch.tensor(served["served_hidden"]) + + def gap(scores): + return float(scores[1] - scores[0]) + + hf_gap = gap(weight @ hf + bias) + vllm_gap = gap(weight @ vllm + bias) + print(f"\n[{name}] HF/vLLM final activation (recorded, not asserted)") + print(f" HF / vLLM norm: {hf.norm():.4f} / {vllm.norm():.4f}") + print( + f" relative L2 / max abs diff: " + f"{(vllm - hf).norm() / hf.norm():.6f} / {(vllm - hf).abs().max():.4f}" + ) + print( + f" reference score gap (fp32 head, HF activation): {gap(reference_scores):.4f}" + ) + print(f" loaded bf16 head, HF activation: {hf_gap:.4f}") + print(f" loaded bf16 head, vLLM activation: {vllm_gap:.4f}") + print(f" activation contribution: {vllm_gap - hf_gap:.4f}") + + +def test_classifier_reads_an_intermediate_layer_correctly(tmp_path_factory): + """A layer-20 slot under eager and FULL, against an fp32 HF layer-20 reference. + + The fixed seed-0 head's reference margin (about 0.07) is below the score + tolerance, so the winning label is inconclusive and this does not identify + layer 20 against nearby layers; it shows the path serves in both modes with + a score gap within tolerance of the layer-20 reference. + """ + read_layer = 20 # a middle layer of the 40-layer base + result = _compose( + tmp_path_factory, ["compose-at-layer", str(read_layer)], "cls_layer_ckpt" + ) + label_token_ids = result["label_token_ids"] + env = { + "CLS_CKPT": result["ckpt"], + "CLS_PIDS": _write_pids(result), + "CLS_LABEL_IDS": json.dumps(label_token_ids), + } + # Serve both modes before asserting so one failure cannot hide the other. + served_by_mode = { + mode: _run_worker(result["outdir"], [mode], env) + for mode in ("run-at-layer", "run-at-layer-full-cudagraph") + } + failures = [] + for mode, served in served_by_mode.items(): + failures += _check_served_scores( + f"{mode} (layer={read_layer})", + served, + result["ground_truth_scores"], + result["ground_truth_label_id"], + label_token_ids, + min_label_margin=1.0, + ) + assert not failures, "\n".join(failures) + + +def test_two_classifier_slots_replay_correctly_in_one_captured_shape(tmp_path_factory): + """Serve a final-layer and a layer-15 slot under eager and FULL. + + Classifier arithmetic and transport are checked against the head's + actual input, recorded in-process at ``compute_logits``: + + - slot 1 (final layer): packed scores equal the loaded bf16 head applied + in fp32 to the served final activation, rounded to bf16 (exact); + - both slots: rewritten label logits equal packed scores (exact), and + label logprob differences equal logit differences up to fp32 + log-softmax rounding. + + Slot 2's input (layer 15) is not observable at ``compute_logits``, so its + source-selection coverage remains the comparison against the layer-15 HF + reference. Slot 1's HF/vLLM activation difference is printed, with no + acceptance threshold. The test does not prove which captured graph replayed. + """ + read_layer = 15 + result = _compose( + tmp_path_factory, ["compose-two-slot", str(read_layer)], "cls_two_slot_ckpt" + ) + label_token_ids = result["label_token_ids"] + env = { + "CLS_CKPT": result["ckpt"], + "CLS_PIDS_1": json.dumps(result["slot1_prompt_ids"]), + "CLS_PIDS_2": json.dumps(result["slot2_prompt_ids"]), + "CLS_LABEL_IDS": json.dumps(label_token_ids), + } + # Serve both modes before asserting so one failure cannot hide another. + served_by_mode = { + mode: _run_worker(result["outdir"], [mode], env) + for mode in ("run-two-slot-eager", "run-two-slot-full-cudagraph") + } + failures = [] + for mode, served in served_by_mode.items(): + slot1, slot2 = served["slot1"], served["slot2"] + failures += _check_transport(f"{mode} slot1", slot1, label_token_ids) + failures += _check_final_layer_arithmetic(f"{mode} slot1", slot1) + _print_final_layer_activation( + f"{mode} slot1", + slot1, + result["slot1_reference_hidden"], + result["slot1_ground_truth_scores"], + ) + if slot1["token_id"] != result["slot1_ground_truth_label_id"]: + failures.append( + f"{mode} slot1: served token {slot1['token_id']} " + f"({slot1['text']!r}) != reference label " + f"{result['slot1_ground_truth_label_id']}" + ) + failures += _check_transport(f"{mode} slot2", slot2, label_token_ids) + failures += _check_served_scores( + f"{mode} slot2 (read_layer={read_layer})", + slot2, + result["slot2_ground_truth_scores"], + result["slot2_ground_truth_label_id"], + label_token_ids, + min_label_margin=None, + ) + assert not failures, "\n".join(failures) diff --git a/tests/vllm/test_classifier_read_points.py b/tests/vllm/test_classifier_read_points.py deleted file mode 100644 index 87dac3e7..00000000 --- a/tests/vllm/test_classifier_read_points.py +++ /dev/null @@ -1,93 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""CPU tests for ``GraniteSwitchModel._classifier_read_points`` (vLLM). - -The read runs outside the compiled region and only needs ``input_ids``, the -final hidden states and ``query_start_loc``, so it is exercised here with a fake -forward context instead of a served model. Serving coverage, including chunked -prefill, lives in ``test_classifier_chunked_prefill.py``. -""" - -from types import SimpleNamespace - -import pytest -import torch - -pytest.importorskip("vllm") - -import vllm.forward_context - -from granite_switch.vllm.granite_switch_model import GraniteSwitchModel - -# Two classifier slots: control id 90 fires slot 2, 91 fires slot 3. -_SLOT_OF = {90: 2, 91: 3} -_MODEL = SimpleNamespace( - config=SimpleNamespace(classifier_control_token_ids=list(_SLOT_OF)) -) - - -def _read(monkeypatch, slices, query_start_loc=None): - """Run the read over one pass made of the given per-request token slices.""" - ids = torch.tensor([t for s in slices for t in s]) - # The switch's index stream: a marker's slot from the marker onward, per request. - indices = [] - for s in slices: - current = 0 - for t in s: - current = _SLOT_OF.get(t, current) - indices.append(current) - classifier_indices = torch.tensor(indices) - hidden = torch.arange(len(ids), dtype=torch.float32).unsqueeze(1) # row i == i - if query_start_loc is None: - bounds = [0] - for s in slices: - bounds.append(bounds[-1] + len(s)) - query_start_loc = bounds - ctx = SimpleNamespace( - attn_metadata=SimpleNamespace(query_start_loc=torch.tensor(query_start_loc)) - ) - monkeypatch.setattr(vllm.forward_context, "get_forward_context", lambda: ctx) - return GraniteSwitchModel._classifier_read_points( - _MODEL, classifier_indices, hidden, ids - ) - - -def test_reports_each_marker_to_its_requests_row(monkeypatch): - rows, hidden, slots, num_reqs = _read(monkeypatch, [[5, 6, 90], [7, 8], [9, 91]]) - # Markers sit at flat 2 and 6, the last rows of requests 0 and 2. - assert rows.tolist() == [0, 2] - assert hidden.squeeze(1).tolist() == [2.0, 6.0] - assert slots.tolist() == [2, 3] - assert num_reqs == 3 - - -def test_pass_without_a_marker_returns_none(monkeypatch): - assert _read(monkeypatch, [[5, 6], [7]]) is None - - -def test_marker_before_the_last_token_raises(monkeypatch): - with pytest.raises(RuntimeError, match=r"request\(s\) \[1\]"): - _read(monkeypatch, [[5, 90], [90, 6, 7]]) - - -def test_two_markers_raise_even_when_one_is_last(monkeypatch): - with pytest.raises(RuntimeError, match="not their last token"): - _read(monkeypatch, [[90, 5, 91]]) - - -def test_empty_slice_does_not_shift_rows(monkeypatch): - # A padded query_start_loc repeats its final offset: the third slice is empty. - rows, _, slots, num_reqs = _read( - monkeypatch, [[5, 90], [6]], query_start_loc=[0, 2, 3, 3] - ) - assert rows.tolist() == [0] - assert slots.tolist() == [2] - assert num_reqs == 3 - - -def test_no_metadata_returns_none(monkeypatch): - ctx = SimpleNamespace(attn_metadata=None) - monkeypatch.setattr(vllm.forward_context, "get_forward_context", lambda: ctx) - out = GraniteSwitchModel._classifier_read_points( - _MODEL, torch.tensor([2]), torch.zeros(1, 1), torch.tensor([90]) - ) - assert out is None diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index 9617ccc7..8fe733c9 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -50,10 +50,17 @@ size. That is the whole artifact — a single linear layer that maps a base-mode logit per label. Train the head however you like (a frozen base model plus a trainable linear layer over the last -hidden state is the usual recipe). At serving time the head reads the final hidden state of the -control token, which the model embeds as the base's role-open marker: `<|start_of_role|>` on -Granite 4.0/4.1, `<|im_start|>` on 4.2. So train on the chat-rendered prompt with no generation -prompt, with that marker appended, and read the hidden state there. +hidden state is the usual recipe). At serving time the head reads the hidden state of the control +token, which the model embeds as the base's role-open marker: `<|start_of_role|>` on Granite +4.0/4.1, `<|im_start|>` on 4.2. So train on the chat-rendered prompt with no generation prompt, +with that marker appended, and read the hidden state there. + +By default that is the **final post-norm** hidden state (HF's `last_hidden_state`). A slot can +instead read an intermediate decoder layer (see `layer` in Step 2); train it on that layer's output. +With HF's `output_hidden_states=True` that is `hidden_states[layer + 1]`, except for the last +decoder layer, whose entry HF replaces with the normalized output; read that one with a forward hook +on `model.model.layers[layer]`. The head must be trained on the same representation it will read at +serving time. Here is a helper that writes a **synthetic** head. @@ -125,6 +132,19 @@ safety: placement), while `kind` is whether the slot is a real LoRA or a classifier head. A plain LoRA entry simply omits `kind`/`labels`. +A classifier entry may also set an optional `layer`, the zero-based decoder layer whose output the +head reads (the residual stream after that layer, before the final norm). Omit it to read the final +post-norm hidden state: + +```yaml +safety_mid: + path: ./safety-mid-head/safety_mid/granite-4.1-3b/lora + type: lora + kind: classifier + labels: [safe, unsafe] + layer: 20 # 0..num_decoder_layers - 1 (granite-4.1-3b has 40); omit for final +``` + You can mix LoRA and classifier entries in one manifest, and give each classifier its own `labels` list of any length: @@ -205,18 +225,21 @@ Compose will: ``` Classifier slots: - safety → labels ['safe', 'unsafe'] → token ids [19193, 39257] - faithfulness → labels ['good', 'bad', 'neutral'] → token ids [19045, 14176, 60668] + safety → labels ['safe', 'unsafe'] → token ids [19193, 39257] @ final norm + faithfulness → labels ['good', 'bad', 'neutral'] → token ids [19045, 14176, 60668] @ final norm ``` +A slot with `layer: 20` prints `@ layer 20` instead. + The resulting `config.json` carries the per-slot classifier metadata as plain lists/ints, so it round-trips through `save_pretrained`/`from_pretrained` with no custom code, and both the HF and -vLLM backends read it identically. Two fields are the stored contract: +vLLM backends read it identically. Three fields are the stored contract: | Field | Meaning | |-------|---------| | `adapter_kinds` | per-slot `"lora"` / `"classifier"`, one entry per adapter — what marks a slot a classifier | | `classifier_label_token_ids` | per-slot list of label token ids, `null` on LoRA slots | +| `classifier_read_layers` | per-slot zero-based decoder layer a classifier reads; `null` for the final post-norm state and on LoRA slots | The rest are **derived** from those in `GraniteSwitchConfig.__init__`, so they appear on the config object but are not independent inputs: @@ -310,10 +333,15 @@ channel. `max_classifier_labels`; a slot with fewer labels leaves its extra rows zero. This keeps the vLLM head a single fused, compile-safe operation regardless of how many labels each slot has — the same way the LoRA bank pads every adapter to `max_lora_rank`. -- **Label-word verdict exit.** The model collapses each classifier request to its final token, - computes the head's per-label logits, blanks the vocabulary, and scatters the verdict onto that - slot's own label token ids (sliced to the slot's real label count, so padded columns are never - emitted). The sampler then emits the label word. +- **Scored at the read layer.** A slot reading an intermediate layer is scored as soon as that + layer runs; final-layer slots are scored after the final norm. The per-label scores and slot id + travel with the hidden states, so vLLM's own row selection delivers them to the logits step, + which also works under CUDA-graph capture. +- **Label-word verdict exit.** At each classifier request's final token the model blanks the + vocabulary and scatters the head's per-label logits onto that slot's own label token ids (sliced + to the slot's real label count, so padded columns are never emitted). The sampler then emits the + label word. + ## Related From ede52f4da3d342ddf3cd28752d41b43339aabb1a Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Tue, 13 Oct 2026 06:26:27 +0000 Subject: [PATCH 10/20] feat(classifier): read classifier layers on SR and across pipeline stages Compose classifier heads alongside Shadow Residual adapters, choosing SR mode from the adapting slots, and keep HF classifier routing through save/reload. SR classifier slots can now read any decoder layer: the adapted stream for active tokens and the base stream otherwise, before final normalization. Marker-only indices and accumulated scores are carried across pipeline ranks, and layer-output selection is shared through DecoderInterface. The packed slot-count guard is removed. Stack the two PP halves in SRDecoderInterface.enter_decoder_stack instead of concatenating them, so vLLM's graph splitter keeps the 2*M row relation; concatenation produced a partition cycle under compile. Test recorder callbacks move to an importable module so PP workers can receive them. Known issue, deferred: under vLLM 0.26, FULL CUDA-graph replay diverges from eager for the standard (non-SR) 3B classifier fixtures. Signed-off-by: IRENEKO --- src/granite_switch/composer/arch.py | 14 +- src/granite_switch/composer/compose_utils.py | 9 +- src/granite_switch/composer/validator.py | 4 +- src/granite_switch/config.py | 15 -- src/granite_switch/hf/core/classifier.py | 7 +- src/granite_switch/hf/switch/multi.py | 8 +- src/granite_switch/vllm/core/classifier.py | 7 +- src/granite_switch/vllm/decoder/interface.py | 28 ++- .../vllm/granite_switch_model.py | 146 +++++---------- tests/composer/test_sr_compose_e2e.py | 109 +++++++++++ tests/shared/classifier_recorder.py | 54 ++++++ tests/unit/test_config.py | 27 ++- .../_classifier_chunked_prefill_worker.py | 174 +++++++++++++----- tests/vllm/test_classifier_capture_safety.py | 174 ++++++++++++++++-- tests/vllm/test_classifier_chunked_prefill.py | 132 +++++++++++++ tutorials/guides/classifier_slots.md | 19 +- 16 files changed, 699 insertions(+), 228 deletions(-) create mode 100644 tests/shared/classifier_recorder.py diff --git a/src/granite_switch/composer/arch.py b/src/granite_switch/composer/arch.py index 279de384..c5b46100 100644 --- a/src/granite_switch/composer/arch.py +++ b/src/granite_switch/composer/arch.py @@ -394,14 +394,6 @@ def granite_dense_arch(base_config=None) -> ArchDescriptor: ) -_SR_BUFFER_KEYWORDS: list[str] = [ - "adapter_token_ids", - "adapter_scalings", - "control_to_substitute_lut", - "cross_stream", # zero-initialized, no base source -] - - def granite_moe_hybrid_sr_arch(base_config=None) -> ArchDescriptor: """GraniteMoeHybrid Shadow Residual architecture. @@ -410,7 +402,7 @@ def granite_moe_hybrid_sr_arch(base_config=None) -> ArchDescriptor: """ arch = granite_moe_hybrid_arch(base_config=base_config) arch.groups = arch.groups + _cross_stream_groups() - arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS) + arch.buffer_keywords.append("cross_stream") # zero-initialized, no base source return arch @@ -418,7 +410,7 @@ def granite_moe_sr_arch(base_config=None) -> ArchDescriptor: """GraniteMoe Shadow Residual architecture (attention + ``cross_stream``).""" arch = granite_moe_arch(base_config=base_config) arch.groups = arch.groups + _cross_stream_groups() - arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS) + arch.buffer_keywords.append("cross_stream") return arch @@ -426,7 +418,7 @@ def granite_dense_sr_arch(base_config=None) -> ArchDescriptor: """Granite dense Shadow Residual architecture (fused + ``cross_stream``).""" arch = granite_dense_arch(base_config=base_config) arch.groups = arch.groups + _cross_stream_groups() - arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS) + arch.buffer_keywords.append("cross_stream") return arch diff --git a/src/granite_switch/composer/compose_utils.py b/src/granite_switch/composer/compose_utils.py index 32e90913..a3bc87b9 100644 --- a/src/granite_switch/composer/compose_utils.py +++ b/src/granite_switch/composer/compose_utils.py @@ -118,12 +118,11 @@ def from_base_and_adapters( print(f"Loading config from {base_model_name_or_path}...") base_config = load_base_config(base_model_name_or_path) - # Classify every adapter as single-stream (LoRA/aLoRA) or dual-stream - # (Shadow Residual). A checkpoint holds one kind or the other: the whole - # decoder runs in one mode, chosen from ``dual_stream``. + # Classify adapting slots as single-stream (LoRA/aLoRA) or dual-stream + # (Shadow Residual). Classifier heads can read either decoder mode. sr_adapters = [] non_sr_adapters = [] - for ap in adapter_paths: + for ap in lora_adapter_paths: if is_shadow_residual_adapter(ap): sr_adapters.append(ap) else: @@ -179,7 +178,7 @@ def from_base_and_adapters( # separately selected LoRA slot can remain active through a classifier # marker because classifier selection does not overwrite adapter state. if lora_adapter_paths: - lora_rank, lora_alpha, lora_ranks, lora_alphas = detect_lora_config( + lora_rank, _lora_alpha, lora_ranks, lora_alphas = detect_lora_config( lora_adapter_paths ) lora_target_modules, source_analysis = detect_present_modules( diff --git a/src/granite_switch/composer/validator.py b/src/granite_switch/composer/validator.py index 56710936..4bba80ff 100644 --- a/src/granite_switch/composer/validator.py +++ b/src/granite_switch/composer/validator.py @@ -65,7 +65,7 @@ def validate_control_lut(model) -> None: def validate_cross_stream_population(model): - """Check that every adapter in a Shadow Residual checkpoint loaded its + """Check that every SR adapter in a Shadow Residual checkpoint loaded its ``cross_stream`` weights. An all-zero slot means the base→adapter injection never fires for that @@ -85,6 +85,8 @@ def validate_cross_stream_population(model): # slots in a freshly built model, but lora_B decides whether the branch # contributes anything. for adapter_idx in range(param.shape[0]): + if model.config.adapter_kinds[adapter_idx] == "classifier": + continue if bool(torch.all(param[adapter_idx] == 0)): empty.append(f"{name}[{adapter_idx}]") diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index 0debdb05..d7e08931 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -441,16 +441,6 @@ def __init__( "verdict as a label word, so the per-label token ids are required. " "Re-compose so config.json carries classifier_label_token_ids." ) - # vLLM packs each token's classifier slot id into the bf16 hidden states, - # which hold integers exactly only up to 256. - if any( - kind == "classifier" and slot > 256 - for slot, kind in enumerate(adapter_kinds, start=1) - ): - raise ValueError( - "classifier slot ids must be <= 256; classifier slots are numbered " - "after LoRA slots, so compose at most 256 adapters." - ) self.classifier_label_token_ids = classifier_label_token_ids # Per slot: zero-based decoder layer a classifier reads (switch cache @@ -481,11 +471,6 @@ def __init__( f"classifier_read_layers[{i}] = {layer!r} must be an int in " f"0..{num_decoder_layers - 1}; omit it for the final layer." ) - if self.dual_stream: - raise ValueError( - f"classifier slot {i} reads layer {layer}, but intermediate " - f"layers are not supported with dual_stream=True." - ) self.classifier_read_layers = classifier_read_layers # Derived: per-slot label counts (0 on LoRA slots) and the padded bank diff --git a/src/granite_switch/hf/core/classifier.py b/src/granite_switch/hf/core/classifier.py index 929bf1d0..5b65d1f5 100644 --- a/src/granite_switch/hf/core/classifier.py +++ b/src/granite_switch/hf/core/classifier.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """Switched classifier heads for Granite Switch. -A classifier head is an optional substitute for a LoRA adapter: an adapter slot -is either a LoRA adapter or a classifier head, selected by the same control -token and adapter index. Where a classifier is selected, the LoRA path no-ops -and the classifier reads clean base-model hidden states. +Classifier markers select a head independently of the active LoRA or SR +adapter. The model supplies the configured layer's residual stream (or final +normalized hidden states), and the head produces per-label scores. Classifier positions are expected to be sparse, and the weight bank is replicated on every rank at this point; revisit if classifiers ever cover most tokens or diff --git a/src/granite_switch/hf/switch/multi.py b/src/granite_switch/hf/switch/multi.py index 549ef415..836224dc 100644 --- a/src/granite_switch/hf/switch/multi.py +++ b/src/granite_switch/hf/switch/multi.py @@ -286,7 +286,7 @@ def __init__( # no slot is a classifier. kind_lut = build_adapter_kind_lut(config) if kind_lut is not None: - self.register_buffer("adapter_kind_lut", kind_lut, persistent=False) + self.register_buffer("adapter_kind_lut", kind_lut, persistent=True) else: self.adapter_kind_lut = None @@ -297,8 +297,10 @@ def __init__( self.control_to_substitute_lut = None marker_lut, slot_lut = build_classifier_control_luts(config) if marker_lut is not None: - self.register_buffer("classifier_control_lut", marker_lut, persistent=False) - self.register_buffer("classifier_slot_lut", slot_lut, persistent=False) + # These routing tables must survive HF's meta-device loading, just + # like the token-exchange LUT above. + self.register_buffer("classifier_control_lut", marker_lut, persistent=True) + self.register_buffer("classifier_slot_lut", slot_lut, persistent=True) else: self.classifier_control_lut = None self.classifier_slot_lut = None diff --git a/src/granite_switch/vllm/core/classifier.py b/src/granite_switch/vllm/core/classifier.py index bf021ad9..32bba285 100644 --- a/src/granite_switch/vllm/core/classifier.py +++ b/src/granite_switch/vllm/core/classifier.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """Switched classifier heads for Granite Switch (vLLM backend). -A classifier head is an optional substitute for a LoRA adapter: an adapter slot -is either a LoRA adapter or a classifier head, selected by the same control -token and adapter index. Where a classifier is selected, the LoRA path no-ops -and the classifier reads clean base-model hidden states. +Classifier markers select a head independently of the active LoRA or SR +adapter. The model supplies the configured layer's residual stream (or final +normalized hidden states), and the head produces per-label scores. Classifier positions are expected to be sparse, and the weight bank is replicated on every rank at this point; revisit if classifiers ever cover most tokens or diff --git a/src/granite_switch/vllm/decoder/interface.py b/src/granite_switch/vllm/decoder/interface.py index adcd324a..433c179b 100644 --- a/src/granite_switch/vllm/decoder/interface.py +++ b/src/granite_switch/vllm/decoder/interface.py @@ -305,6 +305,10 @@ def to_intermediate(self, state) -> tuple[torch.Tensor, torch.Tensor | None]: ``residual=None``) so vLLM's per-token IntermediateTensors slice is safe. """ + @abc.abstractmethod + def layer_output(self, state, adapter_indices) -> torch.Tensor: + """Read this layer's per-token residual stream before the final norm.""" + @abc.abstractmethod def exit_decoder_stack(self, state, adapter_indices, norm, config) -> torch.Tensor: """Last-rank finalize: produce the final ``[M, H]`` hidden states.""" @@ -424,6 +428,11 @@ def run_layer(self, layer, positions, state): def to_intermediate(self, state): return state.hidden_states, state.residual + def layer_output(self, state, adapter_indices): + if state.residual is None: + return state.hidden_states + return state.residual + state.hidden_states + def exit_decoder_stack(self, state, adapter_indices, norm, config): # Fold in the last residual via rms_norm_select so the same fused/separate # convention is used throughout (bit-exact with the original vLLM class). @@ -567,7 +576,10 @@ def enter_decoder_stack(self, hidden_states, residual): if residual is None: hs = torch.cat([hidden_states, hidden_states], dim=0) # [2M, H] else: - hs = torch.cat([hidden_states, residual], dim=0) # [2M, H] + # Both PP halves have M rows. Stack preserves that relation in the + # symbolic shape (2*M), avoiding a late slice-bound dependency when + # vLLM partitions post-attention LoRA metadata reads. + hs = torch.stack([hidden_states, residual], dim=0).flatten(0, 1) return SRStackState(hs=hs) def run_layer(self, layer, positions, state): @@ -580,15 +592,17 @@ def to_intermediate(self, state): m = state.hs.shape[0] // 2 return state.hs[:m].contiguous(), state.hs[m:].contiguous() - def exit_decoder_stack(self, state, adapter_indices, norm, config): + def layer_output(self, state, adapter_indices): hs = state.hs m = adapter_indices.shape[0] - # Per-token merge: adapter-active tokens read the adapter stream, the rest - # read the (base-equivalent) base stream. SR uses plain norm, not - # rms_norm_select (its residual is materialized inside each layer). + # Read after the layer's cross-stream shunt: active tokens use the + # adapted stream, inactive tokens use the base stream. select = (adapter_indices > 0).unsqueeze(1) # [M, 1] - merged = torch.where(select, hs[m:], hs[:m]) - return norm(merged) + return torch.where(select, hs[m:], hs[:m]) + + def exit_decoder_stack(self, state, adapter_indices, norm, config): + # SR materializes residuals inside each layer, so use a plain norm. + return norm(self.layer_output(state, adapter_indices)) # --------------------------------------------------------------------------- # diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index 82d6c618..7fec04de 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -188,13 +188,6 @@ def __init__( if vllm_config.parallel_config.use_ubatching: raise ValueError("Classifier slots do not support microbatching.") - # Built once: copying a list to CUDA inside forward breaks capture. - self.register_buffer( - "classifier_control_token_ids_t", - torch.tensor(config.classifier_control_token_ids, dtype=torch.long), - persistent=False, - ) - is_classifier = [kind == "classifier" for kind in config.adapter_kinds] # Index 0 is a sentinel so a 1-based slot id indexes it directly. self.register_buffer( @@ -223,15 +216,6 @@ def __init__( self._classifier_final_layer_group = name else: self._classifier_layer_groups[layer] = name - # Intermediate scores are not carried across pipeline ranks. - if ( - self._classifier_layer_groups - and vllm_config.parallel_config.pipeline_parallel_size > 1 - ): - raise ValueError( - "Intermediate-layer classifier slots do not support " - "pipeline parallelism." - ) else: self.classifier_head = None self._classifier_final_layer_group = None @@ -315,6 +299,15 @@ def make_empty_intermediate_tensors( device=device, ), } + if self.classifier_head is not None: + tensors["classifier_indices"] = torch.zeros( + batch_size, dtype=torch.long, device=device + ) + tensors["classifier_scores"] = torch.zeros( + (batch_size, self.config.max_classifier_labels), + dtype=dtype, + device=device, + ) return IntermediateTensors(tensors) @@ -386,10 +379,8 @@ def forward( ) modified_input_ids = input_ids - # Split routing state from classifier probe state. In the default - # mode classifier slots are removed from the LoRA stream; markers - # get a separate probe index while the previously selected adapter - # remains active at that token. + # Marker-only probe indices are independent of persistent adapter + # routing. They travel across PP stages along with their scores. if self.switch is not None: classifier_probe_indices = ( self.switch.classifier_indices_from_tokens(input_ids) @@ -403,62 +394,24 @@ def forward( lora_indices = adapter_indices classifier_indices = torch.zeros_like(adapter_indices) - # Prepare kernel metadata ONCE for all decoder layers. The fused - # kernels route on lora_indices, including the marker token when - # coactivation is enabled. - if self.lora_meta is not None and self.lora_ctx is not None: - self.decoder_interface.prepare_kernel_meta( - self.lora_meta, lora_indices, self.lora_ctx - ) - - # Store the persistent adapter state for pipeline parallelism. - # Classifier probe indices are reconstructed from this rank's - # original input_ids and the config-derived marker LUT. - if intermediate_tensors is None: - intermediate_tensors = IntermediateTensors({}) - intermediate_tensors["adapter_indices"] = adapter_indices else: - # Subsequent ranks: recompute fixed-size LoRA metadata from - # token-leading adapter_indices received through PP. - if intermediate_tensors is not None: - adapter_indices = intermediate_tensors["adapter_indices"] - # Recompute both streams from the PP-propagated adapter state - # and this rank's token ids. The classifier marker LUT is a - # non-persistent config-derived buffer on every rank. - if self.switch is not None: - classifier_probe_indices = ( - self.switch.classifier_indices_from_tokens(input_ids) - if input_ids is not None - else None - ) - lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices, classifier_probe_indices - ) - else: - lora_indices = adapter_indices - classifier_indices = torch.zeros_like(adapter_indices) - - if self.lora_ctx is not None: - # Reuse this rank's LoRA stream, including the marker token - # when coactivation is enabled. - self.decoder_interface.prepare_kernel_meta( - self.lora_meta, lora_indices, self.lora_ctx - ) - else: - # Fallback: no metadata available (should not happen in normal operation) - num_tokens = input_ids.shape[0] if input_ids is not None else 0 - if input_ids is not None: - fallback_device = input_ids.device - elif self.lora_meta is not None: - fallback_device = self.lora_meta.device - else: - fallback_device = self.embed_tokens.weight.device - adapter_indices = torch.zeros( - num_tokens, - dtype=torch.long, - device=fallback_device, - ) - classifier_indices = torch.zeros_like(adapter_indices) + assert intermediate_tensors is not None + adapter_indices = intermediate_tensors["adapter_indices"] + # Later stages need no input ids: rank zero already identified the + # markers, without turning them into persistent adapter selection. + classifier_indices = ( + intermediate_tensors["classifier_indices"] + if self.classifier_head is not None + else torch.zeros_like(adapter_indices) + ) + lora_indices = adapter_indices + + # Rebuild this rank's fixed-size kernel metadata from token-leading + # adapter ids; classifier markers preserve the active adapter. + if self.lora_meta is not None and self.lora_ctx is not None: + self.decoder_interface.prepare_kernel_meta( + self.lora_meta, lora_indices, self.lora_ctx + ) # ═══════════════════════════════════════════════════════════════ # Get embeddings (or hidden states from previous pipeline stage) @@ -489,21 +442,23 @@ def forward( # Classifier slots reading an intermediate layer are scored as soon as # that layer runs; only the small score tensor is carried forward. if self.classifier_head is not None: - classifier_scores = torch.zeros( - classifier_indices.shape[0], - self.config.max_classifier_labels, - device=classifier_indices.device, - dtype=hidden_states.dtype, - ) + if get_pp_group().is_first_rank: + classifier_scores = torch.zeros( + classifier_indices.shape[0], + self.config.max_classifier_labels, + device=classifier_indices.device, + dtype=hidden_states.dtype, + ) + else: + classifier_scores = intermediate_tensors["classifier_scores"] state = self.decoder_interface.enter_decoder_stack(hidden_states, residual) for i in range(self.start_layer, self.end_layer): state = self.decoder_interface.run_layer(self.layers[i], positions, state) group = self._classifier_layer_groups.get(i) if group is not None: - # Layer i's output before the final norm: residual + block output. classifier_scores = self._score_group( group, - state.residual + state.hidden_states, + self.decoder_interface.layer_output(state, adapter_indices), classifier_indices, classifier_scores, ) @@ -522,7 +477,7 @@ def forward( # logits_indices gather selects them; compute_logits unpacks them. if self.classifier_head is not None: classifier_scores, classifier_slot = self._compute_classifier_outputs( - input_ids, hidden_states, classifier_indices, classifier_scores + hidden_states, classifier_indices, classifier_scores ) hidden_states = torch.cat( [hidden_states, classifier_scores, classifier_slot.unsqueeze(-1)], @@ -539,6 +494,10 @@ def forward( h_out, r_out = self.decoder_interface.to_intermediate(state) intermediate_tensors["hidden_states"] = h_out intermediate_tensors["residual"] = r_out + intermediate_tensors["adapter_indices"] = adapter_indices + if self.classifier_head is not None: + intermediate_tensors["classifier_indices"] = classifier_indices + intermediate_tensors["classifier_scores"] = classifier_scores return intermediate_tensors def _score_group( @@ -555,26 +514,19 @@ def _score_group( def _compute_classifier_outputs( self, - input_ids: torch.Tensor | None, final_hidden_states: torch.Tensor, classifier_indices: torch.Tensor, classifier_scores: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Add final-layer scores and return per-token (scores, slot ids). - Only marker tokens carry a nonzero slot id. Callers must place one - classifier marker last in each request; misplaced or duplicate markers - are not rejected (request-boundary validation is deferred). + ``classifier_indices`` identifies only this forward's marker tokens, + computed on the first PP rank and passed unchanged between stages. + Callers must place one classifier marker last in each request; + misplaced or duplicate markers are not rejected (request-boundary + validation is deferred). """ - is_marker = ( - torch.isin(input_ids, self.classifier_control_token_ids_t) - if input_ids is not None - else torch.zeros_like(classifier_indices, dtype=torch.bool) - ) - # Exact in bf16: the config limits classifier slot ids to 256. - classifier_slot = torch.where( - is_marker, classifier_indices, torch.zeros_like(classifier_indices) - ).to(final_hidden_states.dtype) + classifier_slot = classifier_indices.to(final_hidden_states.dtype) if self._classifier_final_layer_group is not None: classifier_scores = self._score_group( self._classifier_final_layer_group, diff --git a/tests/composer/test_sr_compose_e2e.py b/tests/composer/test_sr_compose_e2e.py index 7a889b58..90c542f8 100644 --- a/tests/composer/test_sr_compose_e2e.py +++ b/tests/composer/test_sr_compose_e2e.py @@ -184,6 +184,115 @@ def _create_sr_adapter( save_file(state_dict, str(path / "adapter_model.safetensors")) +def _compose_sr_classifier(path: Path, read_layers=(None,)): + """Tiny SR + classifier checkpoint, also used by the vLLM test. + + An identity head exposes every component of its input, making a base-stream + read distinguishable from an adapted-stream read without random margins. + """ + from granite_switch.composer import GraniteSwitchComposer + from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE + + base, sr, classifier = (path / name for name in ("base", "sr", "classifier")) + base.mkdir(parents=True) + classifier.mkdir() + _create_base_model(base) + _create_sr_adapter(sr) + save_file( + {"weight": torch.eye(HIDDEN), "bias": torch.zeros(HIDDEN)}, + str(classifier / CLASSIFIER_HEAD_FILE), + ) + return GraniteSwitchComposer.from_base_and_adapters( + base_model_name_or_path=str(base), + adapter_paths=[str(sr), *[str(classifier)] * len(read_layers)], + adapter_token_ids=list( + range(CONTROL_TOKEN_ID, CONTROL_TOKEN_ID + 1 + len(read_layers)) + ), + adapter_substitute_token_ids=[1, *[2] * len(read_layers)], + adapter_names=["sr", *[f"detect{i}" for i in range(len(read_layers))]], + adapter_kinds=["lora", *["classifier"] * len(read_layers)], + classifier_label_token_ids=[ + None, + *[list(range(100, 100 + HIDDEN))] * len(read_layers), + ], + classifier_read_layers=[None, *read_layers], + switch_head_dim=16, + ).eval() + + +@pytest.mark.parametrize("read_layer", [None, 0, NUM_LAYERS - 1]) +def test_sr_classifier_reads_selected_stream(tmp_path, capsys, read_layer): + """Identity scores prove stream, layer and normalization before/after reload.""" + from granite_switch.hf import load_model + + model = _compose_sr_classifier(tmp_path, read_layers=[read_layer]) + assert "cross_stream slots are all-zero" not in capsys.readouterr().out + assert model.config.dual_stream + head = model.model.classifier_head + torch.testing.assert_close(head.weight[1], torch.eye(HIDDEN)) + assert torch.count_nonzero(head.weight[0]) == 0 + for layer in model.model.layers: + assert torch.count_nonzero(layer.cross_stream.lora_B[0]) > 0 + assert torch.count_nonzero(layer.cross_stream.lora_B[1]) == 0 + + streams = [] + hook = model.model.layers[ + -1 if read_layer is None else read_layer + ].register_forward_hook(lambda _module, _args, output: streams.append(output[0])) + ids = torch.tensor([[10, CONTROL_TOKEN_ID, 20, CONTROL_TOKEN_ID + 1]]) + try: + with torch.no_grad(): + output = model(input_ids=ids, use_cache=False, output_hidden_states=True) + finally: + hook.remove() + assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1]] + assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] + base, adapted = ( + (model.model.norm(h) if read_layer is None else h)[0, -1] for h in streams[0] + ) + assert not torch.allclose(adapted, base), "fixture must exercise SR adaptation" + if read_layer is None: + torch.testing.assert_close(output.hidden_states[-1][0, -1], adapted) + else: + assert not torch.allclose(adapted, output.hidden_states[-1][0, -1]) + labels = model.config.classifier_label_token_ids[1] + torch.testing.assert_close(output.logits[0, -1, labels], adapted) + non_labels = torch.ones(VOCAB_SIZE, dtype=torch.bool) + non_labels[labels] = False + assert torch.isneginf(output.logits[0, -1, non_labels]).all() + + checkpoint = tmp_path / "composed" + model.save_pretrained(checkpoint) + loaded = load_model(str(checkpoint)).eval() + for name in ("adapter_kind_lut", "classifier_control_lut", "classifier_slot_lut"): + torch.testing.assert_close( + getattr(loaded.model.switch, name), getattr(model.model.switch, name) + ) + with torch.no_grad(): + restored = loaded(input_ids=ids, use_cache=False) + torch.testing.assert_close(restored.logits, output.logits) + assert loaded.model._last_lora_indices[0, -1].item() == 1 + assert loaded.model._last_classifier_indices[0, -1].item() == 2 + + # With no SR activation, the probe must read the base stream at that layer. + streams.clear() + hook = model.model.layers[ + -1 if read_layer is None else read_layer + ].register_forward_hook(lambda _module, _args, result: streams.append(result[0])) + try: + with torch.no_grad(): + inactive = model( + input_ids=torch.tensor([[10, 21, 20, CONTROL_TOKEN_ID + 1]]), + use_cache=False, + ) + finally: + hook.remove() + expected = streams[0][0] + if read_layer is None: + expected = model.model.norm(expected) + torch.testing.assert_close(inactive.logits[0, -1, labels], expected[0, -1]) + + # --------------------------------------------------------------------------- # Fixtures: compose once, save to disk, share across all tests # --------------------------------------------------------------------------- diff --git a/tests/shared/classifier_recorder.py b/tests/shared/classifier_recorder.py new file mode 100644 index 00000000..290ea297 --- /dev/null +++ b/tests/shared/classifier_recorder.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Importable model-RPC callbacks for recording sampled classifier rows. + +PP's multiprocessing executor uses standard pickle, so callbacks must live at +module scope. Each worker owns its records; the driver reads them through RPC. +""" + + +def install_classifier_recorder(model): + """Record after compute_logits, outside the compiled/captured forward.""" + from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM + + target = next(m for m in model.modules() if isinstance(m, GraniteSwitchForCausalLM)) + records = [] + target._classifier_test_records = records + original = target.compute_logits + config = target.config + head = target.model.classifier_head + + def compute_logits(hidden_states): + logits = original(hidden_states) + slots = hidden_states[:, -1].round().long() + for row in (slots > 0).nonzero(as_tuple=True)[0].tolist(): + slot = int(slots[row]) + labels = config.classifier_label_token_ids[slot - 1] + n = len(labels) + h = config.hidden_size + record = { + "slot": slot, + "packed_scores": hidden_states[row, h : h + n].float().cpu(), + "rewritten_label_logits": logits[row, labels].float().cpu(), + } + # Only a final-layer head reads the packed hidden-state prefix. + if config.classifier_read_layers[slot - 1] is None: + record["served_hidden"] = hidden_states[row, :h].float().cpu() + record["loaded_weight"] = head.weight[slot - 1, :n].float().cpu() + record["loaded_bias"] = head.bias[slot - 1, :n].float().cpu() + records.append(record) + return logits + + target.compute_logits = compute_logits + + +def read_classifier_records(model): + """Return and clear this worker's rows; non-sampling PP ranks return [].""" + from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM + + target = next(m for m in model.modules() if isinstance(m, GraniteSwitchForCausalLM)) + records = target._classifier_test_records + result = [ + {k: v if k == "slot" else v.tolist() for k, v in r.items()} for r in records + ] + records.clear() + return result diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 2ad4393d..95d51862 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -308,22 +308,21 @@ def test_out_of_range_or_non_int_layer_raises(self, layer): **_classifier_kwargs(classifier_read_layers=[None, layer]) ) - def test_dual_stream_with_intermediate_layer_raises(self): - with pytest.raises(ValueError, match="dual_stream=True"): - GraniteSwitchConfig( - **_classifier_kwargs( - classifier_read_layers=[None, 2], - dual_stream=True, - cross_stream_rank=8, - ) + def test_dual_stream_with_intermediate_layer(self): + cfg = GraniteSwitchConfig( + **_classifier_kwargs( + classifier_read_layers=[None, 2], + dual_stream=True, + cross_stream_rank=8, ) + ) + assert cfg.classifier_read_layers == [None, 2] + restored = GraniteSwitchConfig.from_dict(cfg.to_dict()) + assert restored.dual_stream + assert restored.classifier_read_layers == cfg.classifier_read_layers - def test_classifier_slot_id_above_256_raises(self): - with pytest.raises(ValueError, match="<= 256"): - GraniteSwitchConfig(**_classifier_kwargs(num_adapters=257)) - - def test_classifier_slot_id_256_is_accepted(self): - GraniteSwitchConfig(**_classifier_kwargs(num_adapters=256)) + def test_classifier_slot_id_has_no_config_cap(self): + GraniteSwitchConfig(**_classifier_kwargs(num_adapters=257)) def test_dual_stream_with_final_layer_is_accepted(self): GraniteSwitchConfig(**_classifier_kwargs(dual_stream=True, cross_stream_rank=8)) diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py index d70860b4..d6c8fd45 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -299,51 +299,16 @@ def compose_two_slot(base, outdir, read_layer, control_offset=5, num_labels=2): def _record_classifier_rows(llm): - """Copy each sampled classifier row as ``compute_logits`` sees it. + """Install after startup; PP callbacks must be importable for standard pickle.""" + from tests.shared.classifier_recorder import install_classifier_recorder - Installed after engine startup, so profiling and capture runs are not - recorded. ``compute_logits`` runs outside the captured forward, and the - wrapper only copies tensors after the original returns. Requires the - in-process engine so the returned list is filled in this process. - """ - from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM + llm.apply_model(install_classifier_recorder) - records = [] - def install(model): - target = next( - m for m in model.modules() if isinstance(m, GraniteSwitchForCausalLM) - ) - original = target.compute_logits - config = target.config - head = target.model.classifier_head - - def compute_logits(hidden_states): - logits = original(hidden_states) - slots = hidden_states[:, -1].round().long() - for row in (slots > 0).nonzero(as_tuple=True)[0].tolist(): - slot = int(slots[row]) - labels = config.classifier_label_token_ids[slot - 1] - n = len(labels) - h = config.hidden_size - record = { - "slot": slot, - "packed_scores": hidden_states[row, h : h + n].float().cpu(), - "rewritten_label_logits": logits[row, labels].float().cpu(), - } - # The packed hidden-state prefix is the final post-norm - # activation: the actual head input only for final-layer slots. - if config.classifier_read_layers[slot - 1] is None: - record["served_hidden"] = hidden_states[row, :h].float().cpu() - record["loaded_weight"] = head.weight[slot - 1, :n].float().cpu() - record["loaded_bias"] = head.bias[slot - 1, :n].float().cpu() - records.append(record) - return logits +def _read_classifier_rows(llm): + from tests.shared.classifier_recorder import read_classifier_records - target.compute_logits = compute_logits - - llm.apply_model(install) - return records + return [row for rows in llm.apply_model(read_classifier_records) for row in rows] def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, eager): @@ -369,15 +334,15 @@ def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, if not eager: kw["compilation_config"] = {"cudagraph_mode": "FULL"} llm = LLM(**kw) - records = _record_classifier_rows(llm) + _record_classifier_rows(llm) sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=len(label_token_ids)) results = [] for expected_slot, prompt_ids in ((1, prompt_ids_1), (2, prompt_ids_2)): - records.clear() out = llm.generate([TokensPrompt(prompt_token_ids=prompt_ids)], sp)[0].outputs[ 0 ] + records = _read_classifier_rows(llm) recorded_slots = [r["slot"] for r in records] if recorded_slots != [expected_slot]: raise RuntimeError( @@ -392,7 +357,7 @@ def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, "token_id": int(out.token_ids[0]), "text": out.text, "label_logprobs": label_logprobs, - **{k: v.tolist() for k, v in records[0].items() if k != "slot"}, + **{k: v for k, v in records[0].items() if k != "slot"}, } ) @@ -433,6 +398,90 @@ def run_at_layer(ckpt, prompt_ids, label_token_ids, out_path, full_cudagraph=Fal ) +def run_sr_classifier(ckpt, out_path, *, eager, pp_size=1): + """Exercise SR probes at each configured layer, sequentially and in a mixed batch.""" + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + if pp_size > 1: + # Only trusted, importable test callbacks cross the local RPC boundary. + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + from transformers import AutoConfig + from vllm import LLM, SamplingParams + from vllm.inputs import TokensPrompt + + from tests.composer.test_sr_compose_e2e import CONTROL_TOKEN_ID, HIDDEN + + labels = list(range(100, 100 + HIDDEN)) + config = AutoConfig.from_pretrained(ckpt) + markers = config.classifier_control_token_ids + kw = dict( + model=ckpt, + skip_tokenizer_init=True, + enforce_eager=eager, + dtype="bfloat16", + gpu_memory_utilization=0.3, + max_model_len=128, + max_num_batched_tokens=128, + max_num_seqs=16, + max_logprobs=HIDDEN + 1, + enable_prefix_caching=False, + pipeline_parallel_size=pp_size, + ) + if pp_size > 1: + kw["distributed_executor_backend"] = "mp" + if not eager: + kw["compilation_config"] = { + "cudagraph_mode": "FULL", + "cudagraph_capture_sizes": [1, 2, 4, 16], + } + llm = LLM(**kw) + _record_classifier_rows(llm) + sampling = SamplingParams(max_tokens=1, temperature=0.0, logprobs=HIDDEN + 1) + plain = [10, CONTROL_TOKEN_ID, 20, 30] + + def serve(prompts): + outputs = llm.generate( + [TokensPrompt(prompt_token_ids=p) for p in prompts], sampling + ) + return { + "records": _read_classifier_rows(llm), + "outputs": [ + { + "token_id": int(o.outputs[0].token_ids[0]), + "label_logprobs": [ + o.outputs[0].logprobs[0][tid].logprob + if tid in o.outputs[0].logprobs[0] + else None + for tid in labels + ], + "has_non_label_support": any( + tid not in labels and lp.logprob > float("-inf") + for tid, lp in o.outputs[0].logprobs[0].items() + ), + } + for o in outputs + ], + } + + result = { + "sequential": [ + serve([p]) + for marker in markers + for p in ( + [10, CONTROL_TOKEN_ID, 20, marker], + [10, 21, 20, marker], + [10, CONTROL_TOKEN_ID, 20, marker], + ) + ] + } + mixed = [[10, CONTROL_TOKEN_ID, 20, markers[0]], plain, [10, 21, 20, markers[-1]]] + if len(markers) > 1: + mixed.append([10, CONTROL_TOKEN_ID, 20, markers[1]]) + result["mixed"] = serve(mixed) + with open(out_path, "w") as f: + json.dump(result, f) + + def run_full_cudagraph_replay( ckpt, control_id, label_token_ids, out_path, *, eager=False ): @@ -550,6 +599,27 @@ def run_budget(ckpt, prompt_ids, budget, out_path): def main(): mode = sys.argv[1] + if mode in ("compose-sr-classifier", "compose-sr-classifiers"): + from pathlib import Path + + from tests.composer.test_sr_compose_e2e import _compose_sr_classifier + + read_layers = [0, 1, None] if mode == "compose-sr-classifiers" else [None] + model = _compose_sr_classifier( + Path(os.environ["CLS_OUT"]), read_layers=read_layers + ) + ckpt = os.path.join(os.environ["CLS_OUT"], "composed") + model.save_pretrained(ckpt) + print(json.dumps({"result": {"ckpt": ckpt}})) + return + if mode in ("run-sr-eager", "run-sr-full", "run-sr-pp-eager", "run-sr-pp-full"): + run_sr_classifier( + os.environ["CLS_CKPT"], + sys.argv[3] if mode.startswith("run-sr-pp-") else sys.argv[2], + eager=mode.endswith("eager"), + pp_size=int(sys.argv[2]) if mode.startswith("run-sr-pp-") else 1, + ) + return if mode == "compose-multilabel": try: meta = compose(os.environ["CLS_BASE"], os.environ["CLS_OUT"]) @@ -656,4 +726,18 @@ def main(): if __name__ == "__main__": - main() + if sys.argv[1].startswith("run-sr-pp-"): + # Like _pp_generation_worker, contain distributed-worker teardown in + # this subprocess so a completed test cannot hang in Python finalizers. + exit_code = 1 + try: + main() + exit_code = 0 + except BaseException: + traceback.print_exc() + finally: + sys.stdout.flush() + sys.stderr.flush() + os._exit(exit_code) + else: + main() diff --git a/tests/vllm/test_classifier_capture_safety.py b/tests/vllm/test_classifier_capture_safety.py index c0285dad..00bd2ae2 100644 --- a/tests/vllm/test_classifier_capture_safety.py +++ b/tests/vllm/test_classifier_capture_safety.py @@ -5,11 +5,19 @@ ``test_classifier_chunked_prefill.py``. Request-boundary validation is deferred. """ +from types import SimpleNamespace + import pytest import torch pytest.importorskip("vllm") +from vllm.sequence import IntermediateTensors + +from granite_switch.vllm.decoder.interface import ( + LoRADecoderInterface, + SRDecoderInterface, +) from granite_switch.vllm.granite_switch_model import ( GraniteSwitchForCausalLM, GraniteSwitchModel, @@ -20,7 +28,7 @@ class _StubClassifierHead: """Score = slot id * hidden[:, 0], so a score shows which slot and source it read.""" def __call__(self, x, classifier_indices): - return (classifier_indices.float() * x[:, 0]).unsqueeze(-1).expand(-1, 2) + return (classifier_indices.to(x.dtype) * x[:, 0]).unsqueeze(-1).expand(-1, 2) class _Model(torch.nn.Module): @@ -38,7 +46,6 @@ def __init__(self, slot_of_marker, read_layers=None, kinds=None): is_classifier = [k == "classifier" for k in kinds or ["classifier"] * num_slots] self.config = type("Cfg", (), {"max_classifier_labels": 2})() self.classifier_head = _StubClassifierHead() - self.classifier_control_token_ids_t = torch.tensor(sorted(slot_of_marker)) self.classifier_slot_mask_t = torch.tensor([False, *is_classifier]) slots_by_layer = {} for slot, (is_clf, layer) in enumerate( @@ -57,18 +64,12 @@ def __init__(self, slot_of_marker, read_layers=None, kinds=None): self._classifier_layer_groups[layer] = name -def _run(model, slot_of_marker, tokens, input_ids_override=...): +def _run(model, slot_of_marker, tokens): """Call ``_compute_classifier_outputs`` on a flat token list; hidden[t] = t + 1.""" - indices, current = [], 0 - for t in tokens: - current = slot_of_marker.get(t, current) - indices.append(current) - input_ids = ( - torch.tensor(tokens) if input_ids_override is ... else input_ids_override - ) + indices = [slot_of_marker.get(t, 0) for t in tokens] hidden = torch.arange(1.0, len(tokens) + 1).unsqueeze(1).expand(-1, 4) return GraniteSwitchModel._compute_classifier_outputs( - model, input_ids, hidden, torch.tensor(indices), torch.zeros(len(tokens), 2) + model, hidden, torch.tensor(indices), torch.zeros(len(tokens), 2) ) @@ -110,10 +111,13 @@ def test_no_marker_is_all_zero(): assert not scores.any() -def test_no_input_ids_is_all_zero_slot(): +def test_transferred_marker_indices_need_no_input_ids(): model = _Model({90: 2, 91: 3}) - _, slot = _run(model, {90: 2, 91: 3}, [5, 90, 6], input_ids_override=None) - assert slot.tolist() == [0, 0, 0] + scores, slot = GraniteSwitchModel._compute_classifier_outputs( + model, torch.ones(3, 4), torch.tensor([0, 2, 0]), torch.zeros(3, 2) + ) + assert slot.tolist() == [0, 2, 0] + assert scores.tolist() == [[0.0, 0.0], [2.0, 2.0], [0.0, 0.0]] def test_final_layer_scoring_leaves_intermediate_slot_scores_untouched(): @@ -121,7 +125,6 @@ def test_final_layer_scoring_leaves_intermediate_slot_scores_untouched(): pre_scores = torch.tensor([[0.0, 0.0], [99.0, 99.0], [0.0, 0.0]]) scores, slot = GraniteSwitchModel._compute_classifier_outputs( model, - torch.tensor([5, 90, 91]), torch.ones(3, 4), torch.tensor([0, 2, 3]), pre_scores, @@ -185,10 +188,149 @@ def test_scoring_resolves_replaced_registered_buffers(): scores = _score(model, 0, 3.0, [1, 2], torch.zeros(2, 2)) scores, slots = GraniteSwitchModel._compute_classifier_outputs( model, - torch.tensor([90, 91]), torch.full((2, 4), 5.0), torch.tensor([1, 2]), scores, ) assert scores.tolist() == [[5.0, 5.0], [6.0, 6.0]] assert slots.tolist() == [1.0, 2.0] + + +@pytest.mark.parametrize("sr", [False, True], ids=["lora", "sr"]) +@pytest.mark.parametrize("stages", [2, 3]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_pipeline_handoff_without_later_token_ids(monkeypatch, sr, stages, dtype): + """Run the real forward and decoder strategies, with deterministic layer outputs. + + A layer-0 probe crosses the boundary; another probe and final norm are scored + on the last rank. With three stages, the middle rank only relays scores. + Distinct values expose resets, wrong layers, stream selection and reliance + on input_ids. The token slice matches vLLM's PP buffer contract. + """ + from granite_switch.vllm import granite_switch_model as model_module + + class Switch: + def __call__(self, input_ids, **kwargs): + return (input_ids != 5).long(), input_ids + + def classifier_indices_from_tokens(self, input_ids): + lut = torch.zeros(100, dtype=torch.long) + lut[90:93] = torch.tensor([2, 3, 4]) + return lut[input_ids] + + def split_indices(self, indices, probes): + return indices, probes + + class Layer(torch.nn.Module): + def forward(self, positions, hidden_states, residual=None): + if sr: + base, adapted = hidden_states.chunk(2) + return torch.cat([base + 1, adapted + 10]) + if residual is None: + residual = torch.zeros_like(hidden_states) + return hidden_states + 1, residual + + model = _Model( + {90: 2, 91: 3, 92: 4}, + read_layers=[None, 0, stages - 1, None], + kinds=["lora", "classifier", "classifier", "classifier"], + ) + model.config.hidden_size = 4 + model.config.embedding_multiplier = 1.0 + model.config.fused_add_norm = False + model.switch = Switch() + model.adapter_token_ids = torch.tensor([89, 90, 91, 92]) + model.lora_meta = model.lora_ctx = None + model.decoder_interface = SRDecoderInterface() if sr else LoRADecoderInterface() + model.layers = torch.nn.ModuleList([Layer() for _ in range(stages)]) + model.norm = torch.nn.Identity() + model._compute_classifier_outputs = ( + GraniteSwitchModel._compute_classifier_outputs.__get__(model) + ) + model.start_layer, model.end_layer = 0, 1 + monkeypatch.setattr( + model_module, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=False), + ) + wire = GraniteSwitchModel.forward( + model, + torch.tensor([5, 90, 91, 92, 5]), + torch.arange(5), + inputs_embeds=torch.ones(5, 4, dtype=dtype), + )[1:] + assert isinstance(wire, IntermediateTensors) + assert wire["classifier_indices"].tolist() == [2, 3, 4, 0] + assert wire["classifier_scores"][:, 0].tolist() == [22 if sr else 4, 0, 0, 0] + buffers = GraniteSwitchModel.make_empty_intermediate_tensors( + model, 4, dtype, torch.device("cpu") + ) + assert buffers.tensors.keys() == wire.tensors.keys() + for key, tensor in wire.items(): + assert tensor.shape == buffers[key].shape + assert tensor.dtype == buffers[key].dtype + buffers[key].copy_(tensor) + + for rank in range(1, stages): + model.start_layer, model.end_layer = rank, rank + 1 + group = SimpleNamespace(is_first_rank=False, is_last_rank=rank == stages - 1) + monkeypatch.setattr(model_module, "get_pp_group", lambda group=group: group) + packed = GraniteSwitchModel.forward( + model, None, torch.arange(4), intermediate_tensors=buffers + ) + if not group.is_last_rank: + assert torch.equal(packed["classifier_scores"], wire["classifier_scores"]) + for key, tensor in packed.items(): + buffers[key].copy_(tensor) + assert packed[:, -1].tolist() == [2, 3, 4, 0] + source = 1 + (10 if sr else 1) * stages + expected = [22 if sr else 4, 3 * source, 4 * source, 0] + assert packed[:, 4].tolist() == expected + assert packed[-1, 0] == stages + 1 # inactive SR token reads the base stream + + +def test_sr_pipeline_stack_partitions_post_attention_metadata(): + """A later PP stage must retain 2*M through attention and metadata slicing. + + Concatenating independently symbolic PP inputs produces a late 2*((A+B)//2) + slice bound. vLLM moves that metadata getitem to its producer's partition, + creating a cycle. The real SR entry hook must allow splitting and reuse for + changed token counts, without changing the two streams' values. + """ + from vllm.compilation.backends import split_graph + + @torch.library.custom_op("granite_switch_test::sr_pp_attention", mutates_args=()) + def attention(query: torch.Tensor) -> torch.Tensor: + return query.clone() + + @attention.register_fake + def fake_attention(query): + return torch.empty_like(query) + + interface = SRDecoderInterface() + compilations = [] + + def backend(graph, inputs): + compilations.append(graph) + return split_graph(graph, ["granite_switch_test::sr_pp_attention"])[0].forward + + def stage(base, adapted, indices): + metadata = torch.cat([torch.zeros_like(indices), indices]).unsqueeze(0) + hs = interface.enter_decoder_stack(base, adapted).hs + m = hs.shape[0] // 2 + query = torch.cat([hs[:m], hs[m:]], dim=1) + attended = attention(query) + stacked = torch.cat([attended[:, :4], attended[:, 4:]], dim=0) + remapped = metadata[0, : stacked.shape[0]] + return stacked + remapped.unsqueeze(-1).to(stacked.dtype) + + compiled = torch.compile(stage, backend=backend, dynamic=True, fullgraph=True) + for m in (4, 6): + base = torch.arange(m * 4, dtype=torch.float32).reshape(m, 4) + adapted = base + 17 + indices = torch.arange(1, m + 1) + for tensor in (base, adapted, indices): + torch._dynamo.mark_dynamic(tensor, 0) + expected = torch.cat([base, adapted + indices.unsqueeze(-1)]) + assert torch.equal(compiled(base, adapted, indices), expected) + assert len(compilations) == 1 diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_classifier_chunked_prefill.py index 6c7e3519..aa70b5f9 100644 --- a/tests/vllm/test_classifier_chunked_prefill.py +++ b/tests/vllm/test_classifier_chunked_prefill.py @@ -90,6 +90,11 @@ def _run_worker(outdir, argv, env, timeout=1800): text=True, timeout=timeout, ) + # Keep complete worker traces; pytest's concise failure message has only + # the tail and can otherwise lose the graph-compiler's originating error. + for stream in ("stdout", "stderr"): + with open(out_path.removesuffix(".json") + f".{stream}.log", "w") as log: + log.write(getattr(proc, stream)) label = " ".join(argv) if not os.path.exists(out_path): pytest.fail( @@ -115,6 +120,133 @@ def _composed_env(composed, **extra): return {"CLS_CKPT": composed["ckpt"], "CLS_PIDS": composed["pids_path"], **extra} +@pytest.fixture(scope="module") +def composed_sr(tmp_path_factory): + return _compose(tmp_path_factory, ["compose-sr-classifier"], "sr_cls_ckpt") + + +@pytest.fixture(scope="module") +def composed_sr_probes(tmp_path_factory): + return _compose(tmp_path_factory, ["compose-sr-classifiers"], "sr_probes_ckpt") + + +@pytest.mark.parametrize("mode", ["eager", "full"]) +def test_sr_and_final_classifier_coactivate(composed_sr, mode): + """The identity head scores the selected SR output; transport stays exact.""" + served = _run_worker( + composed_sr["outdir"], + [f"run-sr-{mode}"], + {"CLS_CKPT": composed_sr["ckpt"]}, + ) + sequential = served["sequential"] + assert [len(case["records"]) for case in sequential] == [1, 1, 1] + active, inactive, replay = [case["records"][0] for case in sequential] + assert active["served_hidden"] != inactive["served_hidden"] + assert replay["packed_scores"] == active["packed_scores"] + for case in [*sequential, served["mixed"]]: + for record in case["records"]: + assert record["slot"] == 2 + assert record["packed_scores"] == record["served_hidden"] + assert record["rewritten_label_logits"] == record["packed_scores"] + for case in sequential: + output, record = case["outputs"][0], case["records"][0] + assert not output["has_non_label_support"] + logprobs, scores = output["label_logprobs"], record["packed_scores"] + assert all(lp is not None for lp in logprobs) + assert output["token_id"] == 100 + scores.index(max(scores)) + for lp, score in zip(logprobs, scores): + assert lp - logprobs[0] == pytest.approx(score - scores[0], abs=1e-5) + assert len(served["mixed"]["records"]) == 2 + assert [o["has_non_label_support"] for o in served["mixed"]["outputs"]] == [ + False, + True, + False, + ] + + +@pytest.fixture(scope="module", params=["eager", "full"]) +def sr_probe_reference(composed_sr_probes, request): + mode = request.param + return mode, _run_worker( + composed_sr_probes["outdir"], + [f"run-sr-pp-{mode}", "1"], + {"CLS_CKPT": composed_sr_probes["ckpt"]}, + ) + + +def _check_sr_layer_probes(result, mode, pp): + """Distinct identity-head rows expose wrong sources, masks and stale replay.""" + failures = [] + assert len(result["sequential"]) == 9 + active_scores = [] + for slot in (2, 3, 4): + active, inactive, replay = result["sequential"][3 * (slot - 2) : 3 * (slot - 1)] + assert [r["slot"] for r in active["records"]] == [slot] + assert [r["slot"] for r in inactive["records"]] == [slot] + scores = active["records"][0]["packed_scores"] + active_scores.append(scores) + assert scores != inactive["records"][0]["packed_scores"] + assert scores == replay["records"][0]["packed_scores"] + for name, case in ( + ("active", active), + ("inactive", inactive), + ("replay", replay), + ): + record, output = case["records"][0], case["outputs"][0] + failures.extend( + _check_transport( + f"PP{pp}/{mode}/slot{slot}/{name}", + {**record, **output}, + list(range(100, 164)), + ) + ) + assert output["token_id"] == 100 + record["packed_scores"].index( + max(record["packed_scores"]) + ) + assert not output["has_non_label_support"] + assert all( + active_scores[a] != active_scores[b] for a, b in ((0, 1), (0, 2), (1, 2)) + ) + assert sorted(r["slot"] for r in result["mixed"]["records"]) == [2, 3, 4] + assert [o["has_non_label_support"] for o in result["mixed"]["outputs"]] == [ + False, + True, + False, + False, + ] + assert not failures, "\n".join(failures) + + +def test_sr_classifier_reads_each_layer_under_replay(sr_probe_reference): + """One GPU: probes at layer 0, layer 1 and final norm with SR active/inactive.""" + mode, served = sr_probe_reference + _check_sr_layer_probes(served, mode, 1) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs for PP=2") +def test_sr_layer_probes_across_pipeline_stages(composed_sr_probes, sr_probe_reference): + """Same checkpoint and execution mode: PP=2 must preserve PP=1 score rows.""" + mode, reference = sr_probe_reference + served = _run_worker( + composed_sr_probes["outdir"], + [f"run-sr-pp-{mode}", "2"], + {"CLS_CKPT": composed_sr_probes["ckpt"]}, + ) + _check_sr_layer_probes(served, mode, 2) + # Compare execution modes separately: graph compilation can change bf16 + # rounding, whereas the pipeline handoff itself must preserve score rows. + for case1, case2 in zip( + [*reference["sequential"], reference["mixed"]], + [*served["sequential"], served["mixed"]], + ): + assert case1["outputs"] == case2["outputs"] + assert len(case1["records"]) == len(case2["records"]) + for r1, r2 in zip(case1["records"], case2["records"]): + assert r1["slot"] == r2["slot"] + assert r1["packed_scores"] == r2["packed_scores"] + assert r1["rewritten_label_logits"] == r2["rewritten_label_logits"] + + def test_full_cudagraph_mode_does_not_crash_and_matches_eager(composed): """Exercise FULL startup and changed inputs, including mixed requests. diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index 8fe733c9..09979e8b 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -9,8 +9,8 @@ serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detecti A Granite Switch adapter slot is normally a LoRA adapter. A **classifier slot** is an alternative occupant of that same slot: instead of a LoRA, it holds a small linear classification head. It is triggered by the *same* control-token machinery — one `<|name|>` token, one adapter index — but -where a LoRA would modify the hidden states, a classifier reads the clean base-model hidden state -and produces per-label logits. +the classifier reads the configured layer's activation and produces per-label logits. An active +LoRA or SR adapter stays active while the classifier reads its resulting hidden state. At inference the verdict does not come back as a raw score vector. Each label word is resolved to a single token id at compose time, and the model rewrites the request's final logit row so that the @@ -142,9 +142,14 @@ safety_mid: type: lora kind: classifier labels: [safe, unsafe] - layer: 20 # 0..num_decoder_layers - 1 (granite-4.1-3b has 40); omit for final + layer: 20 # 0..num_decoder_layers - 1; omit for final ``` +This also works with Shadow Residual (SR) adapters: an active adapter's tokens read the adapted +stream after that layer's cross-stream shunt; inactive tokens read the base stream. Decoder layer +numbers exclude the switch's KV-cache slots. Pipeline parallelism preserves the same numbering +across stages. + You can mix LoRA and classifier entries in one manifest, and give each classifier its own `labels` list of any length: @@ -326,9 +331,8 @@ channel. ## How it works (brief) -- **Same trigger as a LoRA.** One `<|name|>` control token per slot, one adapter index. The switch - tags that index as a classifier and routes the token out of the LoRA stream, so the LoRA path - no-ops there and the classifier head runs on the clean base hidden state. +- **Independent probe trigger.** One `<|name|>` control token selects a classifier head while + preserving the active LoRA or SR adapter. The head reads the resulting activation at that token. - **Padded head bank.** All classifier heads live in one stacked tensor padded to `max_classifier_labels`; a slot with fewer labels leaves its extra rows zero. This keeps the vLLM head a single fused, compile-safe operation regardless of how many labels each slot has — the same @@ -337,6 +341,9 @@ channel. layer runs; final-layer slots are scored after the final norm. The per-label scores and slot id travel with the hidden states, so vLLM's own row selection delivers them to the logits step, which also works under CUDA-graph capture. +- **Pipeline transport.** The first stage identifies marker tokens. Each stage scores its own + configured layers and passes those scores and marker indices to the next stage; intermediate + hidden states do not need to be saved for classification on the last stage. - **Label-word verdict exit.** At each classifier request's final token the model blanks the vocabulary and scatters the head's per-label logits onto that slot's own label token ids (sliced to the slot's real label count, so padded columns are never emitted). The sampler then emits the From ddaf4c430789aa3e448dd027f628140a3c727f77 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Tue, 13 Oct 2026 06:29:00 +0000 Subject: [PATCH 11/20] feat(classifier): select an adapter and a classifier independently Chat templates now accept classifier_name alongside adapter_name: the adapter activates at its usual location and the classifier marker is placed last. adapter_name naming a classifier keeps the classifier-only behavior. SR and assistant-boundary aLoRA keep the activation generation prompt; LoRA and user-invoked aLoRA use the transcript-only read point. Existing checkpoints need regenerated templates, not new weights. Signed-off-by: IRENEKO --- .../composer/tokenizer_setup.py | 49 ++++-- tests/composer/test_chat_template.py | 159 ++++++++++++++++++ tutorials/guides/classifier_slots.md | 29 +++- 3 files changed, 220 insertions(+), 17 deletions(-) diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 7dc83dee..6eb20839 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -907,7 +907,10 @@ def configure_chat_template( Modifies the tokenizer's chat template so that callers can pass ``adapter_name="..."`` to ``apply_chat_template()`` and have the - correct control token inserted automatically: + correct control token inserted automatically. An independent + ``classifier_name="..."`` appends a classifier marker as the final token, + preserving the selected adapter. ``adapter_name`` alone still accepts a + classifier slot for compatibility. * **LoRA** adapters: token at the **beginning** of the sequence. * **ALoRA** adapters: token immediately before ``alora_invocation_tokens`` @@ -915,6 +918,10 @@ def configure_chat_template( requirement-checker), or right before the generation prompt for adapters whose invocation sequence is the assistant role token sequence and therefore does not appear in any user message. + * **SR** adapters: replace the generation prompt's activation anchor. + Classification preserves this prompt, as it does an ALoRA fallback + invocation. LoRA and user-invoked ALoRA classification omit the generation + prompt, matching classifier-only placement after the last turn's close. ALoRA placement uses a two-pass Jinja2 approach embedded in the template: @@ -1046,7 +1053,7 @@ def configure_chat_template( {%- set adapter_type = '' %} {%- set adapter_invocation_text = '' %} {%- set adapter_invocation_drop_chars = 1 %} -{%- if adapter_name is defined and adapter_name in adapter_map %} +{%- if adapter_name is defined and adapter_name is string 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' %} @@ -1054,11 +1061,26 @@ def configure_chat_template( {%- set adapter_invocation_drop_chars = adapter_map[adapter_name]['invocation_drop_chars'] %} {%- endif %} {%- endif %} -{#- A classifier reads its verdict at the marker, so nothing may follow it: - drop the generation prompt even when the caller asks for one. -#} -{%- set classifier_mode = adapter_type == 'classifier' %} +{#- Classifier selection is independent of persistent adapter selection. + adapter_name still accepts a classifier for existing callers. -#} +{%- set classifier_token = adapter_token if adapter_type == 'classifier' else '' %} +{%- if classifier_name is defined and classifier_name is not none %} +{%- if classifier_name is not string or classifier_name not in adapter_map or adapter_map[classifier_name]['type'] != 'classifier' %} +{{- raise_exception('classifier_name must name a composed classifier slot') }} +{%- endif %} +{%- if classifier_token and classifier_token != adapter_map[classifier_name]['token'] %} +{{- raise_exception('adapter_name and classifier_name select different classifiers') }} +{%- endif %} +{%- set classifier_token = adapter_map[classifier_name]['token'] %} +{%- endif %} +{%- set classifier_mode = classifier_token != '' %} {%- if classifier_mode %} -{%- set add_generation_prompt = false %} +{%- if adapter_name is defined and adapter_name is not none and not adapter_type %} +{{- raise_exception('adapter_name must name a composed adapter or classifier slot') }} +{%- endif %} +{#- SR needs its generation-prompt anchor before the classifier marker. + The ALoRA fallback below similarly restores its full invocation. -#} +{%- set add_generation_prompt = adapter_type == 'sr' %} {%- endif %} """ @@ -1185,6 +1207,7 @@ def configure_chat_template( (alora_target_idx == -1), meaning the adapter activates at the assistant role token boundary rather than inside a user message. -#} {%- if ns.adapter_token and ns.adapter_type == 'alora' and ns.alora_target_idx == -1 %} +{%- if classifier_mode %}{%- set add_generation_prompt = true %}{%- endif %} {{- ns.adapter_token }} {%- set ns.skip_next_role_marker = true %} {%- endif %} @@ -1331,13 +1354,6 @@ def configure_chat_template( + modified_chat_template[insertion_point:] ) - # Classifier placement: the marker is the very last token, after the last - # turn's close. - if any(i["type"] == ANCHOR_MODE_CLASSIFIER for i in adapter_mapping.values()): - modified_chat_template += ( - "\n{%- if classifier_mode %}{{- adapter_token }}{%- endif %}" - ) - # Insert alora fallback before generation prompt gen_prompt_pattern = r"(\{%- if add_generation_prompt %\})" match = re.search(gen_prompt_pattern, modified_chat_template) @@ -1438,6 +1454,13 @@ def _split_merged(match: "re.Match") -> str: modified_chat_template, ) + # Append only after all activation insertions, including the fallback used + # by templates without a generation-prompt block. Nothing follows the probe. + if any(i["type"] == ANCHOR_MODE_CLASSIFIER for i in adapter_mapping.values()): + modified_chat_template += ( + "\n{%- if classifier_mode %}{{- classifier_token }}{%- endif %}" + ) + tokenizer.chat_template = modified_chat_template print(f"Chat template configured with {len(adapter_mapping)} adapter mappings:") for adapter_name, info in adapter_mapping.items(): diff --git a/tests/composer/test_chat_template.py b/tests/composer/test_chat_template.py index 5cbabc0c..9bbeb63c 100644 --- a/tests/composer/test_chat_template.py +++ b/tests/composer/test_chat_template.py @@ -23,6 +23,7 @@ - End-to-end: adapter_config.json → render (no patching) """ +import json import os from types import SimpleNamespace from unittest.mock import patch @@ -1423,3 +1424,161 @@ def test_sr_anchor_id_must_match_the_tokenizer(self): configure_chat_template( tokenizer, [(self._SR_GRANITE, "answerability", "sr")] ) + + +class TestClassifierSelection: + """The saved template accepts independent selections through the real HF API.""" + + @pytest.fixture( + params=[_GRANITE_TEMPLATE, _CHATML_TEMPLATE], ids=["granite", "chatml"] + ) + def tokenizer(self, request): + from tokenizers import Tokenizer, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + backend = Tokenizer( + models.WordLevel({"[UNK]": 0, "assistant": 1}, unk_token="[UNK]") + ) + backend.pre_tokenizer = pre_tokenizers.Whitespace() + return PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token="[UNK]", + additional_special_tokens=[*_SRTokenizer._SPECIALS, ""], + chat_template=request.param, + ) + + def _configure(self, tokenizer, tmp_path, kind="lora"): + from granite_switch.composer.tokenizer_setup import add_control_tokens + + fmt = detect_template_format(tokenizer.chat_template) + adapter = tmp_path / "rag" + adapter.mkdir() + if kind == "sr": + anchor = "" if fmt.name == "chatml" else "<|end_of_role|>" + config = { + "last_context_token": anchor, + "last_context_token_id": tokenizer.convert_tokens_to_ids(anchor), + } + else: + invocation = ( + "" + if kind == "alora" + else fmt.role_open_marker + + "assistant" + + ("\n" if fmt.name == "chatml" else "<|end_of_role|>") + ) + config = { + "alora_invocation_tokens": tokenizer.encode( + invocation, add_special_tokens=False + ) + } + (adapter / "adapter_config.json").write_text(json.dumps(config)) + discovered = [ + (str(adapter), "rag", "alora" if kind == "alora_fallback" else kind), + (None, "safety", "classifier"), + (None, "other", "classifier"), + ] + add_control_tokens(tokenizer, discovered) + configure_chat_template(tokenizer, discovered) + + _MESSAGES = [ + {"role": "system", "content": "Check this."}, + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "Now Hello"}, + ] + + def test_classifier_only_and_legacy_are_identical(self, tokenizer, tmp_path): + self._configure(tokenizer, tmp_path) + plain = tokenizer.apply_chat_template( + self._MESSAGES, tokenize=False, add_generation_prompt=False + ) + for kwargs in ( + {"classifier_name": "safety"}, + {"adapter_name": "safety"}, + {"adapter_name": "safety", "classifier_name": "safety"}, + ): + result = tokenizer.apply_chat_template( + self._MESSAGES, tokenize=False, add_generation_prompt=True, **kwargs + ) + assert result == plain + "<|safety|>" + + @pytest.mark.parametrize("kind", ["lora", "alora", "alora_fallback", "sr"]) + @pytest.mark.parametrize("gen", [False, True]) + def test_adapter_activation_precedes_final_probe( + self, tokenizer, tmp_path, kind, gen + ): + self._configure(tokenizer, tmp_path, kind) + for thinking in (False, True): + # SR and assistant-boundary aLoRA need the generation-prompt anchor. + # Other cases preserve the transcript-only classifier read point. + expected = tokenizer.apply_chat_template( + self._MESSAGES, + tokenize=False, + adapter_name="rag", + add_generation_prompt=kind in ("sr", "alora_fallback"), + enable_thinking=thinking, + ) + kwargs = dict( + adapter_name="rag", + classifier_name="safety", + add_generation_prompt=gen, + enable_thinking=thinking, + ) + result = tokenizer.apply_chat_template( + self._MESSAGES, tokenize=False, **kwargs + ) + ids = tokenizer.apply_chat_template( + self._MESSAGES, tokenize=True, return_dict=True, **kwargs + )["input_ids"] + assert result == expected + "<|safety|>" + assert ids[-1] == tokenizer.convert_tokens_to_ids("<|safety|>") + assert ids.count(ids[-1]) == 1 + assert ids.count(tokenizer.convert_tokens_to_ids("<|rag|>")) == 1 + + @pytest.mark.parametrize( + "kwargs,match", + [ + ({"classifier_name": "missing"}, "classifier_name"), + ({"classifier_name": "rag"}, "classifier_name"), + ({"classifier_name": ["safety"]}, "classifier_name"), + ( + {"adapter_name": "safety", "classifier_name": "other"}, + "different classifiers", + ), + ({"adapter_name": "missing", "classifier_name": "safety"}, "adapter_name"), + ], + ) + def test_invalid_selection_raises(self, tokenizer, tmp_path, kwargs, match): + from jinja2 import TemplateError + + self._configure(tokenizer, tmp_path) + with pytest.raises(TemplateError, match=match): + tokenizer.apply_chat_template(self._MESSAGES, tokenize=False, **kwargs) + + def test_no_probe_preserves_adapter_render(self, tokenizer, tmp_path): + self._configure(tokenizer, tmp_path) + for gen in (False, True): + kwargs = dict(adapter_name="rag", add_generation_prompt=gen, tokenize=False) + assert tokenizer.apply_chat_template( + self._MESSAGES, classifier_name=None, **kwargs + ) == (tokenizer.apply_chat_template(self._MESSAGES, **kwargs)) + + def test_saved_tokenizer_keeps_both_selections(self, tokenizer, tmp_path): + from transformers import PreTrainedTokenizerFast + + self._configure(tokenizer, tmp_path) + kwargs = dict( + adapter_name="rag", classifier_name="safety", add_generation_prompt=True + ) + expected = tokenizer.apply_chat_template( + self._MESSAGES, return_dict=True, **kwargs + )["input_ids"] + tokenizer.save_pretrained(tmp_path / "tokenizer") + restored = PreTrainedTokenizerFast.from_pretrained(tmp_path / "tokenizer") + assert ( + restored.apply_chat_template(self._MESSAGES, return_dict=True, **kwargs)[ + "input_ids" + ] + == expected + ) diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index 09979e8b..5fb271b9 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -295,9 +295,9 @@ two labels for a 2-label slot, three for a 3-label slot. ### Firing via the chat template (optional) -Instead of placing `<|safety|>` by hand, you can let the composed model's chat template place it for -you by passing `adapter_name` — a classifier slot is activated exactly like any other adapter. The -template appends the control token after the last turn: +Instead of placing `<|safety|>` by hand, pass `classifier_name` to the composed model's chat +template. It appends the control token as the final token. Existing calls that select a classifier +through `adapter_name` still work: ```bash curl http://localhost:8000/v1/chat/completions \ @@ -306,10 +306,31 @@ curl http://localhost:8000/v1/chat/completions \ "model": "./granite-switch-detect", "messages": [{"role": "user", "content": "Is this text safe? The weather is nice today."}], "max_tokens": 1, - "chat_template_kwargs": {"adapter_name": "safety"} + "chat_template_kwargs": {"classifier_name": "safety"} }' ``` +To apply an adapter and classify its resulting activation in the same request, select both: + +```json +"chat_template_kwargs": {"adapter_name": "rag", "classifier_name": "safety"} +``` + +The adapter stays active at the classifier marker. LoRA and user-invoked aLoRA classification omit +the assistant generation prompt. SR and assistant-boundary aLoRA preserve the generation prompt +needed to activate the adapter, then append the classifier marker. This placement is enforced even +when the caller supplies the opposite `add_generation_prompt` value; ChatML thinking branches keep +their usual suffix before the marker. Return the verdict with `max_tokens: 1`. + +Unknown classifier names, non-classifier slots supplied as `classifier_name`, unknown adapters +paired with classification, and conflicting classifier selections raise template errors. +Omitting `classifier_name` or setting it to `null` preserves ordinary adapter-only rendering. + +These arguments are embedded in newly composed tokenizers. Existing checkpoints need their chat +template regenerated (or re-composition) to use the new interface. No model-weight change is needed. +The request-level arguments select one adapter and one classifier; per-message adapter switching +is not exposed by this interface. + ### Reading logprobs (optional) To see the raw verdict rather than just the winning word, ask for a couple of logprobs: From 5f85c95d5ce579348b0eddebe45a15484065c4c9 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 19:48:59 +0000 Subject: [PATCH 12/20] test(classifier): expect the default read layer in manifest metadata Classifier manifest entries without a layer report layer None. Signed-off-by: IRENEKO --- tests/composer/test_adapter_loader.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/composer/test_adapter_loader.py b/tests/composer/test_adapter_loader.py index 372506d0..e27e316f 100644 --- a/tests/composer/test_adapter_loader.py +++ b/tests/composer/test_adapter_loader.py @@ -578,6 +578,7 @@ def test_kind_and_labels_are_parsed(self, tmp_path): "hallucination": { "kind": "classifier", "labels": ["faithful", "hallucinated"], + "layer": None, } } From a134f671e25c648b180fe274409e30c814b23574 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 19:49:03 +0000 Subject: [PATCH 13/20] revert(classifier): drop the independent classifier_name selection adapter_name selects one slot, a generation adapter or a classifier. Joint classifier/adapter selection is out of scope. Signed-off-by: IRENEKO --- .../composer/tokenizer_setup.py | 49 ++---- tests/composer/test_chat_template.py | 159 ------------------ tutorials/guides/classifier_slots.md | 29 +--- 3 files changed, 17 insertions(+), 220 deletions(-) diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 6eb20839..7dc83dee 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -907,10 +907,7 @@ def configure_chat_template( Modifies the tokenizer's chat template so that callers can pass ``adapter_name="..."`` to ``apply_chat_template()`` and have the - correct control token inserted automatically. An independent - ``classifier_name="..."`` appends a classifier marker as the final token, - preserving the selected adapter. ``adapter_name`` alone still accepts a - classifier slot for compatibility. + correct control token inserted automatically: * **LoRA** adapters: token at the **beginning** of the sequence. * **ALoRA** adapters: token immediately before ``alora_invocation_tokens`` @@ -918,10 +915,6 @@ def configure_chat_template( requirement-checker), or right before the generation prompt for adapters whose invocation sequence is the assistant role token sequence and therefore does not appear in any user message. - * **SR** adapters: replace the generation prompt's activation anchor. - Classification preserves this prompt, as it does an ALoRA fallback - invocation. LoRA and user-invoked ALoRA classification omit the generation - prompt, matching classifier-only placement after the last turn's close. ALoRA placement uses a two-pass Jinja2 approach embedded in the template: @@ -1053,7 +1046,7 @@ def configure_chat_template( {%- set adapter_type = '' %} {%- set adapter_invocation_text = '' %} {%- set adapter_invocation_drop_chars = 1 %} -{%- if adapter_name is defined and adapter_name is string and adapter_name in adapter_map %} +{%- 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' %} @@ -1061,26 +1054,11 @@ def configure_chat_template( {%- set adapter_invocation_drop_chars = adapter_map[adapter_name]['invocation_drop_chars'] %} {%- endif %} {%- endif %} -{#- Classifier selection is independent of persistent adapter selection. - adapter_name still accepts a classifier for existing callers. -#} -{%- set classifier_token = adapter_token if adapter_type == 'classifier' else '' %} -{%- if classifier_name is defined and classifier_name is not none %} -{%- if classifier_name is not string or classifier_name not in adapter_map or adapter_map[classifier_name]['type'] != 'classifier' %} -{{- raise_exception('classifier_name must name a composed classifier slot') }} -{%- endif %} -{%- if classifier_token and classifier_token != adapter_map[classifier_name]['token'] %} -{{- raise_exception('adapter_name and classifier_name select different classifiers') }} -{%- endif %} -{%- set classifier_token = adapter_map[classifier_name]['token'] %} -{%- endif %} -{%- set classifier_mode = classifier_token != '' %} +{#- A classifier reads its verdict at the marker, so nothing may follow it: + drop the generation prompt even when the caller asks for one. -#} +{%- set classifier_mode = adapter_type == 'classifier' %} {%- if classifier_mode %} -{%- if adapter_name is defined and adapter_name is not none and not adapter_type %} -{{- raise_exception('adapter_name must name a composed adapter or classifier slot') }} -{%- endif %} -{#- SR needs its generation-prompt anchor before the classifier marker. - The ALoRA fallback below similarly restores its full invocation. -#} -{%- set add_generation_prompt = adapter_type == 'sr' %} +{%- set add_generation_prompt = false %} {%- endif %} """ @@ -1207,7 +1185,6 @@ def configure_chat_template( (alora_target_idx == -1), meaning the adapter activates at the assistant role token boundary rather than inside a user message. -#} {%- if ns.adapter_token and ns.adapter_type == 'alora' and ns.alora_target_idx == -1 %} -{%- if classifier_mode %}{%- set add_generation_prompt = true %}{%- endif %} {{- ns.adapter_token }} {%- set ns.skip_next_role_marker = true %} {%- endif %} @@ -1354,6 +1331,13 @@ def configure_chat_template( + modified_chat_template[insertion_point:] ) + # Classifier placement: the marker is the very last token, after the last + # turn's close. + if any(i["type"] == ANCHOR_MODE_CLASSIFIER for i in adapter_mapping.values()): + modified_chat_template += ( + "\n{%- if classifier_mode %}{{- adapter_token }}{%- endif %}" + ) + # Insert alora fallback before generation prompt gen_prompt_pattern = r"(\{%- if add_generation_prompt %\})" match = re.search(gen_prompt_pattern, modified_chat_template) @@ -1454,13 +1438,6 @@ def _split_merged(match: "re.Match") -> str: modified_chat_template, ) - # Append only after all activation insertions, including the fallback used - # by templates without a generation-prompt block. Nothing follows the probe. - if any(i["type"] == ANCHOR_MODE_CLASSIFIER for i in adapter_mapping.values()): - modified_chat_template += ( - "\n{%- if classifier_mode %}{{- classifier_token }}{%- endif %}" - ) - tokenizer.chat_template = modified_chat_template print(f"Chat template configured with {len(adapter_mapping)} adapter mappings:") for adapter_name, info in adapter_mapping.items(): diff --git a/tests/composer/test_chat_template.py b/tests/composer/test_chat_template.py index 9bbeb63c..5cbabc0c 100644 --- a/tests/composer/test_chat_template.py +++ b/tests/composer/test_chat_template.py @@ -23,7 +23,6 @@ - End-to-end: adapter_config.json → render (no patching) """ -import json import os from types import SimpleNamespace from unittest.mock import patch @@ -1424,161 +1423,3 @@ def test_sr_anchor_id_must_match_the_tokenizer(self): configure_chat_template( tokenizer, [(self._SR_GRANITE, "answerability", "sr")] ) - - -class TestClassifierSelection: - """The saved template accepts independent selections through the real HF API.""" - - @pytest.fixture( - params=[_GRANITE_TEMPLATE, _CHATML_TEMPLATE], ids=["granite", "chatml"] - ) - def tokenizer(self, request): - from tokenizers import Tokenizer, models, pre_tokenizers - from transformers import PreTrainedTokenizerFast - - backend = Tokenizer( - models.WordLevel({"[UNK]": 0, "assistant": 1}, unk_token="[UNK]") - ) - backend.pre_tokenizer = pre_tokenizers.Whitespace() - return PreTrainedTokenizerFast( - tokenizer_object=backend, - unk_token="[UNK]", - additional_special_tokens=[*_SRTokenizer._SPECIALS, ""], - chat_template=request.param, - ) - - def _configure(self, tokenizer, tmp_path, kind="lora"): - from granite_switch.composer.tokenizer_setup import add_control_tokens - - fmt = detect_template_format(tokenizer.chat_template) - adapter = tmp_path / "rag" - adapter.mkdir() - if kind == "sr": - anchor = "" if fmt.name == "chatml" else "<|end_of_role|>" - config = { - "last_context_token": anchor, - "last_context_token_id": tokenizer.convert_tokens_to_ids(anchor), - } - else: - invocation = ( - "" - if kind == "alora" - else fmt.role_open_marker - + "assistant" - + ("\n" if fmt.name == "chatml" else "<|end_of_role|>") - ) - config = { - "alora_invocation_tokens": tokenizer.encode( - invocation, add_special_tokens=False - ) - } - (adapter / "adapter_config.json").write_text(json.dumps(config)) - discovered = [ - (str(adapter), "rag", "alora" if kind == "alora_fallback" else kind), - (None, "safety", "classifier"), - (None, "other", "classifier"), - ] - add_control_tokens(tokenizer, discovered) - configure_chat_template(tokenizer, discovered) - - _MESSAGES = [ - {"role": "system", "content": "Check this."}, - {"role": "user", "content": "first"}, - {"role": "assistant", "content": "reply"}, - {"role": "user", "content": "Now Hello"}, - ] - - def test_classifier_only_and_legacy_are_identical(self, tokenizer, tmp_path): - self._configure(tokenizer, tmp_path) - plain = tokenizer.apply_chat_template( - self._MESSAGES, tokenize=False, add_generation_prompt=False - ) - for kwargs in ( - {"classifier_name": "safety"}, - {"adapter_name": "safety"}, - {"adapter_name": "safety", "classifier_name": "safety"}, - ): - result = tokenizer.apply_chat_template( - self._MESSAGES, tokenize=False, add_generation_prompt=True, **kwargs - ) - assert result == plain + "<|safety|>" - - @pytest.mark.parametrize("kind", ["lora", "alora", "alora_fallback", "sr"]) - @pytest.mark.parametrize("gen", [False, True]) - def test_adapter_activation_precedes_final_probe( - self, tokenizer, tmp_path, kind, gen - ): - self._configure(tokenizer, tmp_path, kind) - for thinking in (False, True): - # SR and assistant-boundary aLoRA need the generation-prompt anchor. - # Other cases preserve the transcript-only classifier read point. - expected = tokenizer.apply_chat_template( - self._MESSAGES, - tokenize=False, - adapter_name="rag", - add_generation_prompt=kind in ("sr", "alora_fallback"), - enable_thinking=thinking, - ) - kwargs = dict( - adapter_name="rag", - classifier_name="safety", - add_generation_prompt=gen, - enable_thinking=thinking, - ) - result = tokenizer.apply_chat_template( - self._MESSAGES, tokenize=False, **kwargs - ) - ids = tokenizer.apply_chat_template( - self._MESSAGES, tokenize=True, return_dict=True, **kwargs - )["input_ids"] - assert result == expected + "<|safety|>" - assert ids[-1] == tokenizer.convert_tokens_to_ids("<|safety|>") - assert ids.count(ids[-1]) == 1 - assert ids.count(tokenizer.convert_tokens_to_ids("<|rag|>")) == 1 - - @pytest.mark.parametrize( - "kwargs,match", - [ - ({"classifier_name": "missing"}, "classifier_name"), - ({"classifier_name": "rag"}, "classifier_name"), - ({"classifier_name": ["safety"]}, "classifier_name"), - ( - {"adapter_name": "safety", "classifier_name": "other"}, - "different classifiers", - ), - ({"adapter_name": "missing", "classifier_name": "safety"}, "adapter_name"), - ], - ) - def test_invalid_selection_raises(self, tokenizer, tmp_path, kwargs, match): - from jinja2 import TemplateError - - self._configure(tokenizer, tmp_path) - with pytest.raises(TemplateError, match=match): - tokenizer.apply_chat_template(self._MESSAGES, tokenize=False, **kwargs) - - def test_no_probe_preserves_adapter_render(self, tokenizer, tmp_path): - self._configure(tokenizer, tmp_path) - for gen in (False, True): - kwargs = dict(adapter_name="rag", add_generation_prompt=gen, tokenize=False) - assert tokenizer.apply_chat_template( - self._MESSAGES, classifier_name=None, **kwargs - ) == (tokenizer.apply_chat_template(self._MESSAGES, **kwargs)) - - def test_saved_tokenizer_keeps_both_selections(self, tokenizer, tmp_path): - from transformers import PreTrainedTokenizerFast - - self._configure(tokenizer, tmp_path) - kwargs = dict( - adapter_name="rag", classifier_name="safety", add_generation_prompt=True - ) - expected = tokenizer.apply_chat_template( - self._MESSAGES, return_dict=True, **kwargs - )["input_ids"] - tokenizer.save_pretrained(tmp_path / "tokenizer") - restored = PreTrainedTokenizerFast.from_pretrained(tmp_path / "tokenizer") - assert ( - restored.apply_chat_template(self._MESSAGES, return_dict=True, **kwargs)[ - "input_ids" - ] - == expected - ) diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index 5fb271b9..09979e8b 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -295,9 +295,9 @@ two labels for a 2-label slot, three for a 3-label slot. ### Firing via the chat template (optional) -Instead of placing `<|safety|>` by hand, pass `classifier_name` to the composed model's chat -template. It appends the control token as the final token. Existing calls that select a classifier -through `adapter_name` still work: +Instead of placing `<|safety|>` by hand, you can let the composed model's chat template place it for +you by passing `adapter_name` — a classifier slot is activated exactly like any other adapter. The +template appends the control token after the last turn: ```bash curl http://localhost:8000/v1/chat/completions \ @@ -306,31 +306,10 @@ curl http://localhost:8000/v1/chat/completions \ "model": "./granite-switch-detect", "messages": [{"role": "user", "content": "Is this text safe? The weather is nice today."}], "max_tokens": 1, - "chat_template_kwargs": {"classifier_name": "safety"} + "chat_template_kwargs": {"adapter_name": "safety"} }' ``` -To apply an adapter and classify its resulting activation in the same request, select both: - -```json -"chat_template_kwargs": {"adapter_name": "rag", "classifier_name": "safety"} -``` - -The adapter stays active at the classifier marker. LoRA and user-invoked aLoRA classification omit -the assistant generation prompt. SR and assistant-boundary aLoRA preserve the generation prompt -needed to activate the adapter, then append the classifier marker. This placement is enforced even -when the caller supplies the opposite `add_generation_prompt` value; ChatML thinking branches keep -their usual suffix before the marker. Return the verdict with `max_tokens: 1`. - -Unknown classifier names, non-classifier slots supplied as `classifier_name`, unknown adapters -paired with classification, and conflicting classifier selections raise template errors. -Omitting `classifier_name` or setting it to `null` preserves ordinary adapter-only rendering. - -These arguments are embedded in newly composed tokenizers. Existing checkpoints need their chat -template regenerated (or re-composition) to use the new interface. No model-weight change is needed. -The request-level arguments select one adapter and one classifier; per-message adapter switching -is not exposed by this interface. - ### Reading logprobs (optional) To see the raw verdict rather than just the winning word, ask for a couple of logprobs: From f9e41bf4d6bff95ae7fcbf5a2018e70887631684 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 20:59:34 +0000 Subject: [PATCH 14/20] refactor(classifier): restore single-slot classifier routing A classifier marker again selects its slot like any adapter: the kind LUT routes it out of the LoRA stream, so no adapter stays active at the marker. This drops the marker-token LUTs that let a classifier coexist with an active adapter, which also removes the vocabulary-sized buffers that went stale after the audio resize and broke strict reloads. The surviving routing helpers move from token_exchange.py into classifier_routing.py. vLLM still masks classifier indices to marker tokens and sends them through PP with the scores. Tests now select classifier and adapter slots separately, and the guide states that a request selects one slot. Signed-off-by: IRENEKO --- src/granite_switch/classifier_routing.py | 58 +++++++++++ src/granite_switch/config.py | 12 +-- src/granite_switch/hf/core/classifier.py | 8 +- .../hf/modeling_granite_switch.py | 30 +++--- src/granite_switch/hf/switch/multi.py | 58 +++-------- src/granite_switch/token_exchange.py | 83 ---------------- src/granite_switch/vllm/core/classifier.py | 8 +- src/granite_switch/vllm/decoder/interface.py | 5 +- .../vllm/granite_switch_model.py | 66 +++++++------ src/granite_switch/vllm/switch/multi.py | 53 +++------- tests/composer/test_classifier_compose.py | 23 +++-- tests/composer/test_control_lut_refresh.py | 41 +++++++- tests/composer/test_sr_compose_e2e.py | 96 +++++++++---------- tests/hf/test_classifier.py | 34 ++++--- tests/unit/test_classifier_routing.py | 72 ++++++++++++++ tests/unit/test_token_exchange.py | 43 --------- .../_classifier_chunked_prefill_worker.py | 21 ++-- tests/vllm/_model_forward_tests.py | 16 ++-- tests/vllm/test_classifier_capture_safety.py | 41 ++++---- tests/vllm/test_classifier_chunked_prefill.py | 40 ++++---- tutorials/guides/classifier_slots.md | 20 ++-- 21 files changed, 416 insertions(+), 412 deletions(-) create mode 100644 src/granite_switch/classifier_routing.py create mode 100644 tests/unit/test_classifier_routing.py diff --git a/src/granite_switch/classifier_routing.py b/src/granite_switch/classifier_routing.py new file mode 100644 index 00000000..5467291c --- /dev/null +++ b/src/granite_switch/classifier_routing.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Classifier-slot routing helpers, shared by both backends' switches. + +An adapter slot is either a generation adapter or a classifier head, selected +by the same control token and adapter index. These helpers split the switch's +single per-token index stream so that classifier positions no-op the LoRA path +and feed the classifier head instead. + +Backend-neutral on purpose: the table is a pure function of +``config.adapter_kinds``, so both switch engines (HF and vLLM MultiSwitch) +build the identical table. +""" + +import torch + + +def build_adapter_kind_lut(config) -> torch.Tensor | None: + """Bool LUT over adapter indices marking which slots are classifiers. + + ``[num_adapters + 1]``, indexed by ADAPTER INDEX: entry ``i`` is True when + adapter index ``i`` is a classifier slot (index 0, base, is always False). + ``None`` when no slot is a classifier, which lets the switch skip the split + entirely. Built from ``config.adapter_kinds``, so both backends' switches + agree on the table by construction. + """ + kinds = getattr(config, "adapter_kinds", None) if config is not None else None + if kinds is None or not any(k == "classifier" for k in kinds): + return None + lut = torch.zeros(len(kinds) + 1, dtype=torch.bool) + for slot, kind in enumerate(kinds): # slot is the 0-based adapter position + if kind == "classifier": + lut[slot + 1] = True # adapter index is slot + 1 + return lut + + +def split_adapter_indices( + kind_lut: torch.Tensor | None, adapter_indices: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Split a per-token index stream into LoRA and classifier streams. + + Shapes are preserved (``[batch, seq_len]`` on HF, ``[total_tokens]`` on + vLLM), so this serves both backends: + + - ``lora_indices``: classifier positions zeroed, so the LoRA path no-ops + there. + - ``classifier_indices``: the adapter index at classifier positions, else 0. + + With ``kind_lut`` None (no classifier slots) the LoRA stream passes through + and the classifier stream is all zeros. + """ + if kind_lut is None: + return adapter_indices, torch.zeros_like(adapter_indices) + is_classifier = kind_lut[adapter_indices] + zeros = torch.zeros_like(adapter_indices) + return ( + torch.where(is_classifier, zeros, adapter_indices), + torch.where(is_classifier, adapter_indices, zeros), + ) diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index d7e08931..6de3b2ca 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -31,9 +31,9 @@ class GraniteSwitchConfig(GraniteMoeHybridConfig): classifier). Default: 0. Index 0 always means "base / no adapter". adapter_token_ids (List[int]): Token IDs for adapter control. Length: num_adapters (or num_adapters + 1 with a leading base-reset - token). Must be unique. A LoRA control token selects that adapter; - a classifier control token selects its head without changing the - active adapter. Output 0 = base. + token). Must be unique. adapter_token_ids[i] activates slot i+1 + (1-indexed output), a LoRA adapter or a classifier head. + Output 0 = base. adapter_substitute_token_ids (List[int]): Token IDs whose embeddings replace the control-token embeddings before the decoder runs. Length: num_adapters (or num_adapters + 1 with a leading base-reset @@ -387,9 +387,9 @@ def __init__( self.adapter_ranks = adapter_ranks # Each adapter slot is either a LoRA adapter ("lora", the default) or - # a classifier head ("classifier"). Classifier markers are orthogonal - # probes: they select a head without replacing the currently active - # adapter selection. The probe slot is looked up from the marker token. + # a classifier head ("classifier"). A classifier slot still uses a + # control token and adapter index; the switch's adapter_kind_lut routes + # its positions out of the LoRA stream and into the classifier head. _VALID_KINDS = ("lora", "classifier") if adapter_kinds is None: adapter_kinds = ["lora"] * num_adapters diff --git a/src/granite_switch/hf/core/classifier.py b/src/granite_switch/hf/core/classifier.py index 5b65d1f5..c23d97b6 100644 --- a/src/granite_switch/hf/core/classifier.py +++ b/src/granite_switch/hf/core/classifier.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Switched classifier heads for Granite Switch. -Classifier markers select a head independently of the active LoRA or SR -adapter. The model supplies the configured layer's residual stream (or final -normalized hidden states), and the head produces per-label scores. +A classifier head is an optional substitute for a LoRA adapter: an adapter slot +is either a LoRA adapter or a classifier head, selected by the same control +token and adapter index. Where a classifier is selected, the LoRA path no-ops. +The model supplies the configured layer's hidden states (the final normalized +state by default), and the head produces per-label scores. Classifier positions are expected to be sparse, and the weight bank is replicated on every rank at this point; revisit if classifiers ever cover most tokens or diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index 0729b26a..a4d812ad 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -419,10 +419,13 @@ def __init__(self, config: GraniteSwitchConfig): self.switch = None self.adapter_token_ids = None - # Classifier heads share the adapter control-token machinery. The marker - # selects a classifier head independently; the active adapter selection - # remains in effect through that token. Intermediate read layers are - # scored as the decoder runs, retaining only marker-row verdicts. + # Classifier heads: optional alternative for LoRA adapters. A classifier + # slot shares the adapter index/token machinery, but its index is tagged + # as "classifier" (via the switch's adapter_kind_lut). The switch's + # split_indices() routes those positions out of the LoRA stream (so LoRA + # no-ops there) and into a classifier stream. The head reads its + # configured layer (the final post-norm state by default), scoring + # intermediate layers as the decoder runs and keeping only marker rows. if "classifier" in getattr(config, "adapter_kinds", []): self.classifier_head = SwitchedClassifierHead( hidden_size=config.hidden_size, @@ -602,17 +605,12 @@ def forward( device=device, ) - # Keep persistent adapter routing separate from the classifier probe: - # classifier markers select a head from their token id and do not - # replace the currently active adapter. + # Split the single adapter_indices stream into a LoRA stream and + # a classifier stream. Classifier positions are zeroed in lora_indices + # (so the LoRA path no-ops there). if self.switch is not None: - classifier_probe_indices = ( - self.switch.classifier_indices_from_tokens(input_ids) - if input_ids is not None - else None - ) lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices, classifier_probe_indices + adapter_indices ) else: lora_indices = adapter_indices @@ -625,8 +623,10 @@ def forward( inputs_embeds = self.embed_tokens(modified_input_ids) inputs_embeds = inputs_embeds * self.embedding_multiplier - # Expose routing for tests/debugging: adapter selection is persistent, - # while classifier indices come independently from marker token ids. + # Expose adapter_indices for tests and debugging. _last_adapter_indices is + # the full stream (LoRA + classifier); _last_lora_indices is what the + # decoder routes on; _last_classifier_indices is the classifier-only + # stream the verdict reads. self._last_adapter_indices = adapter_indices self._last_lora_indices = lora_indices self._last_classifier_indices = classifier_indices diff --git a/src/granite_switch/hf/switch/multi.py b/src/granite_switch/hf/switch/multi.py index 836224dc..3f8f48ae 100644 --- a/src/granite_switch/hf/switch/multi.py +++ b/src/granite_switch/hf/switch/multi.py @@ -66,13 +66,8 @@ from transformers.cache_utils import Cache from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from ...token_exchange import ( - apply_token_exchange, - build_adapter_kind_lut, - build_classifier_control_luts, - build_control_to_substitute_lut, - split_adapter_indices, -) +from ...classifier_routing import build_adapter_kind_lut, split_adapter_indices +from ...token_exchange import apply_token_exchange, build_control_to_substitute_lut from .codes import KerdockDGCodeGenerator, recover_count_from_signal # Large negative finite value for masking. NOT literal -inf because IEEE 754 @@ -283,7 +278,8 @@ def __init__( # rewrite EVERY token id to 0 in apply_token_exchange. # adapter_kind_lut: marks which adapter indices are classifier slots, so # split_indices can route those positions out of the LoRA stream. None when - # no slot is a classifier. + # no slot is a classifier. persistent=True so it survives HF's + # meta-device loading, like the token-exchange LUT below. kind_lut = build_adapter_kind_lut(config) if kind_lut is not None: self.register_buffer("adapter_kind_lut", kind_lut, persistent=True) @@ -295,15 +291,6 @@ def __init__( self.register_buffer("control_to_substitute_lut", lut, persistent=True) else: self.control_to_substitute_lut = None - marker_lut, slot_lut = build_classifier_control_luts(config) - if marker_lut is not None: - # These routing tables must survive HF's meta-device loading, just - # like the token-exchange LUT above. - self.register_buffer("classifier_control_lut", marker_lut, persistent=True) - self.register_buffer("classifier_slot_lut", slot_lut, persistent=True) - else: - self.classifier_control_lut = None - self.classifier_slot_lut = None def _apply(self, *args, **kwargs): """Keep ``codebook`` in fp32 across ``.to()`` / ``.half()`` / ``.bfloat16()``. @@ -333,26 +320,13 @@ def _apply(self, *args, **kwargs): return out def split_indices( - self, - adapter_indices: torch.Tensor, - classifier_probe_indices: torch.Tensor | None = None, + self, adapter_indices: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Split ``[batch, seq_len]`` or ``[num_tokens]`` into (lora, classifier). - See :func:`granite_switch.token_exchange.split_adapter_indices`. + See :func:`granite_switch.classifier_routing.split_adapter_indices`. """ - return split_adapter_indices( - self.adapter_kind_lut, - adapter_indices, - classifier_probe_indices, - ) - - def classifier_indices_from_tokens( - self, input_ids: torch.Tensor - ) -> torch.Tensor | None: - if self.classifier_slot_lut is None: - return None - return self.classifier_slot_lut[input_ids] + return split_adapter_indices(self.adapter_kind_lut, adapter_indices) @property def num_cache_layers(self) -> int: @@ -407,15 +381,9 @@ def forward( # [B, S, 1] == [1, 1, A] -> [B, S, A] matches = input_ids.unsqueeze(2) == adapter_token_ids.unsqueeze(0).unsqueeze(0) is_control_token = matches.any(dim=2) # [B, S] - route_control_token = is_control_token - if self.classifier_control_lut is not None: - # A classifier marker probes the current residual stream; it must - # not overwrite the persistent adapter selection in coded memory. - is_classifier_marker = self.classifier_control_lut[input_ids] - route_control_token = is_control_token & ~is_classifier_marker # expert_id = argmax + offset (offset selects the layout; see __init__). expert_ids = torch.where( - route_control_token, + is_control_token, matches.long().argmax(dim=2) + self._expert_id_offset, torch.zeros_like(input_ids, dtype=torch.long), ) # [B, S] @@ -430,10 +398,8 @@ def forward( device=device, dtype=torch.float32, ) - # Un-mask at the anchor and adapter-selection controls. Classifier - # markers are excluded because they do not change persistent adapter - # selection. - anchor_or_control = is_counting_anchor | route_control_token # [B, S] + # Un-mask (0) dim 0 at anchor + control tokens (arithmetic where). + anchor_or_control = is_counting_anchor | is_control_token # [B, S] key_states_count[:, 0, :, 0] = torch.where( anchor_or_control, torch.zeros_like(key_states_count[:, 0, :, 0]), @@ -559,7 +525,7 @@ def _causal_mask(kv_len: int) -> torch.Tensor: bsz, q_len, self.memory_dim ) # [B, S, memory_dim], fp32 - write_mask = route_control_token.unsqueeze(2).to(torch.float32) # [B, S, 1] + write_mask = is_control_token.unsqueeze(2).to(torch.float32) # [B, S, 1] key_states_memory = torch.zeros( (bsz, 1, q_len, self.memory_head_dim), @@ -578,7 +544,7 @@ def _causal_mask(kv_len: int) -> torch.Tensor: ) value_states_memory[:, 0, :, 0] = expert_ids.to( torch.float32 - ) * route_control_token.to(torch.float32) + ) * is_control_token.to(torch.float32) if past_key_values is not None: cache_kwargs = {"cache_position": cache_position} diff --git a/src/granite_switch/token_exchange.py b/src/granite_switch/token_exchange.py index d73d1aa6..231a2154 100644 --- a/src/granite_switch/token_exchange.py +++ b/src/granite_switch/token_exchange.py @@ -154,86 +154,3 @@ def apply_token_exchange( sub_id_per_pos = lut[input_ids] is_control = sub_id_per_pos >= 0 return torch.where(is_control, sub_id_per_pos, input_ids) - - -def build_adapter_kind_lut(config) -> torch.Tensor | None: - """Bool LUT over adapter indices marking which slots are classifiers. - - ``[num_adapters + 1]``, indexed by ADAPTER INDEX: entry ``i`` is True when - adapter index ``i`` is a classifier slot (index 0, base, is always False). - ``None`` when no slot is a classifier. This supports the legacy split when - no token-derived classifier probe indices are supplied. Built from - ``config.adapter_kinds``, so both backends agree on the table. - """ - kinds = getattr(config, "adapter_kinds", None) if config is not None else None - if kinds is None or not any(k == "classifier" for k in kinds): - return None - lut = torch.zeros(len(kinds) + 1, dtype=torch.bool) - for slot, kind in enumerate(kinds): # slot is the 0-based adapter position - if kind == "classifier": - lut[slot + 1] = True # adapter index is slot + 1 - return lut - - -def build_classifier_control_luts( - config, -) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """Build token-id LUTs for classifier markers. - - The first LUT marks classifier control tokens; the second returns the - classifier head slot (1-based) for each token. These LUTs are separate - from ``adapter_kind_lut`` because an orthogonal classifier probe must not - replace the persistent LoRA adapter selected earlier in the request. - """ - kinds = getattr(config, "adapter_kinds", None) - control_ids = getattr(config, "adapter_token_ids", None) - if not kinds or not control_ids or not any(k == "classifier" for k in kinds): - return None, None - offset = len(control_ids) - len(kinds) - if offset not in (0, 1): - raise ValueError( - "adapter_token_ids must have one id per adapter, optionally with " - "a leading base-reset token" - ) - vocab_size = max( - int(getattr(config, "vocab_size", 0)), - max(control_ids) + 1, - ) - is_marker = torch.zeros(vocab_size, dtype=torch.bool) - slot_lut = torch.zeros(vocab_size, dtype=torch.long) - for slot, kind in enumerate(kinds, start=1): - if kind != "classifier": - continue - token_id = control_ids[slot - 1 + offset] - is_marker[token_id] = True - slot_lut[token_id] = slot - return is_marker, slot_lut - - -def split_adapter_indices( - kind_lut: torch.Tensor | None, - adapter_indices: torch.Tensor, - classifier_probe_indices: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Split a per-token index stream into LoRA and classifier streams. - - Shapes are preserved (``[batch, seq_len]`` on HF, ``[total_tokens]`` on - vLLM), so this serves both backends: - - - When classifier marker indices are provided, they select the classifier - head independently and the active adapter stream is preserved. - - Without marker indices, the legacy kind LUT split is used. - - With ``kind_lut`` None (no classifier slots) the LoRA stream passes through - and the classifier stream is all zeros. - """ - if classifier_probe_indices is not None: - return adapter_indices, classifier_probe_indices - if kind_lut is None: - return adapter_indices, torch.zeros_like(adapter_indices) - is_classifier = kind_lut[adapter_indices] - zeros = torch.zeros_like(adapter_indices) - return ( - torch.where(is_classifier, zeros, adapter_indices), - torch.where(is_classifier, adapter_indices, zeros), - ) diff --git a/src/granite_switch/vllm/core/classifier.py b/src/granite_switch/vllm/core/classifier.py index 32bba285..c3e018ec 100644 --- a/src/granite_switch/vllm/core/classifier.py +++ b/src/granite_switch/vllm/core/classifier.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Switched classifier heads for Granite Switch (vLLM backend). -Classifier markers select a head independently of the active LoRA or SR -adapter. The model supplies the configured layer's residual stream (or final -normalized hidden states), and the head produces per-label scores. +A classifier head is an optional substitute for a LoRA adapter: an adapter slot +is either a LoRA adapter or a classifier head, selected by the same control +token and adapter index. Where a classifier is selected, the LoRA path no-ops. +The model supplies the configured layer's hidden states (the final normalized +state by default), and the head produces per-label scores. Classifier positions are expected to be sparse, and the weight bank is replicated on every rank at this point; revisit if classifiers ever cover most tokens or diff --git a/src/granite_switch/vllm/decoder/interface.py b/src/granite_switch/vllm/decoder/interface.py index 433c179b..81fe9c44 100644 --- a/src/granite_switch/vllm/decoder/interface.py +++ b/src/granite_switch/vllm/decoder/interface.py @@ -595,8 +595,9 @@ def to_intermediate(self, state): def layer_output(self, state, adapter_indices): hs = state.hs m = adapter_indices.shape[0] - # Read after the layer's cross-stream shunt: active tokens use the - # adapted stream, inactive tokens use the base stream. + # Read after the layer's cross-stream shunt: adapter-active tokens use + # the adapted stream, the rest (including classifier markers, whose LoRA + # stream is zeroed) use the base stream. select = (adapter_indices > 0).unsqueeze(1) # [M, 1] return torch.where(select, hs[m:], hs[:m]) diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index 7fec04de..eeba6c81 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -168,9 +168,12 @@ def __init__( self.lora_meta = None self.lora_ctx = None - # Classifier markers share the control-token machinery but select the - # classifier head independently, preserving the active adapter. The head - # scores each slot at its read layer (the final post-norm state by default). + # Classifier heads: optional alternative for LoRA adapters. A classifier + # slot shares the adapter index/token machinery, but its index is tagged + # as "classifier" (via the switch's adapter_kind_lut). The switch's + # split_indices() routes those positions out of the LoRA stream (so LoRA + # no-ops there) and into a classifier stream. The head scores each slot + # at its read layer (the final post-norm state by default). if "classifier" in getattr(config, "adapter_kinds", []): from .core import SwitchedClassifierHead @@ -188,6 +191,13 @@ def __init__( if vllm_config.parallel_config.use_ubatching: raise ValueError("Classifier slots do not support microbatching.") + # Built once: copying a list to CUDA inside forward breaks capture. + self.register_buffer( + "classifier_control_token_ids_t", + torch.tensor(config.classifier_control_token_ids, dtype=torch.long), + persistent=False, + ) + is_classifier = [kind == "classifier" for kind in config.adapter_kinds] # Index 0 is a sentinel so a 1-based slot id indexes it directly. self.register_buffer( @@ -379,35 +389,35 @@ def forward( ) modified_input_ids = input_ids - # Marker-only probe indices are independent of persistent adapter - # routing. They travel across PP stages along with their scores. - if self.switch is not None: - classifier_probe_indices = ( - self.switch.classifier_indices_from_tokens(input_ids) - if input_ids is not None - else None - ) - lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices, classifier_probe_indices - ) - else: - lora_indices = adapter_indices - classifier_indices = torch.zeros_like(adapter_indices) - else: assert intermediate_tensors is not None adapter_indices = intermediate_tensors["adapter_indices"] - # Later stages need no input ids: rank zero already identified the - # markers, without turning them into persistent adapter selection. - classifier_indices = ( - intermediate_tensors["classifier_indices"] - if self.classifier_head is not None - else torch.zeros_like(adapter_indices) + + # Split the single adapter_indices stream into a LoRA stream and a + # classifier stream. Classifier positions are zeroed in lora_indices (so + # the LoRA path no-ops there). + if self.switch is not None: + lora_indices, classifier_indices = self.switch.split_indices( + adapter_indices ) + else: lora_indices = adapter_indices + classifier_indices = torch.zeros_like(adapter_indices) + + # The verdict is read only at the marker itself, not at tokens that + # inherit its slot afterwards. Later PP ranks have no input ids, so they + # receive the marker-only indices from the first rank. + if self.classifier_head is not None: + if not get_pp_group().is_first_rank: + classifier_indices = intermediate_tensors["classifier_indices"] + elif input_ids is not None: + is_marker = torch.isin(input_ids, self.classifier_control_token_ids_t) + classifier_indices = torch.where( + is_marker, classifier_indices, torch.zeros_like(classifier_indices) + ) # Rebuild this rank's fixed-size kernel metadata from token-leading - # adapter ids; classifier markers preserve the active adapter. + # adapter ids. The fused kernel routes on the LoRA stream. if self.lora_meta is not None and self.lora_ctx is not None: self.decoder_interface.prepare_kernel_meta( self.lora_meta, lora_indices, self.lora_ctx @@ -458,7 +468,7 @@ def forward( if group is not None: classifier_scores = self._score_group( group, - self.decoder_interface.layer_output(state, adapter_indices), + self.decoder_interface.layer_output(state, lora_indices), classifier_indices, classifier_scores, ) @@ -470,7 +480,7 @@ def forward( # [M, H] tensor is the final post-norm hidden state (what the LM head # consumes). hidden_states = self.decoder_interface.exit_decoder_stack( - state, adapter_indices, self.norm, self.config + state, lora_indices, self.norm, self.config ) # Pack scores and slot ids onto the hidden states so vLLM's own @@ -520,7 +530,7 @@ def _compute_classifier_outputs( ) -> tuple[torch.Tensor, torch.Tensor]: """Add final-layer scores and return per-token (scores, slot ids). - ``classifier_indices`` identifies only this forward's marker tokens, + ``classifier_indices`` is nonzero only at this forward's marker tokens, computed on the first PP rank and passed unchanged between stages. Callers must place one classifier marker last in each request; misplaced or duplicate markers are not rejected (request-boundary diff --git a/src/granite_switch/vllm/switch/multi.py b/src/granite_switch/vllm/switch/multi.py index 3e0a1460..7c7995f6 100644 --- a/src/granite_switch/vllm/switch/multi.py +++ b/src/granite_switch/vllm/switch/multi.py @@ -98,16 +98,17 @@ from vllm.config import VllmConfig from vllm.model_executor.layers.attention.attention import Attention +from granite_switch.classifier_routing import ( + build_adapter_kind_lut, + split_adapter_indices, +) from granite_switch.hf.switch.codes import ( KerdockDGCodeGenerator, recover_count_from_signal, ) from granite_switch.token_exchange import ( apply_token_exchange, - build_adapter_kind_lut, - build_classifier_control_luts, build_control_to_substitute_lut, - split_adapter_indices, ) # Large negative finite value for masking. NOT literal -inf because IEEE 754 @@ -329,35 +330,15 @@ def __init__( self.register_buffer("control_to_substitute_lut", lut, persistent=True) else: self.control_to_substitute_lut = None - marker_lut, slot_lut = build_classifier_control_luts(hf_config) - if marker_lut is not None: - self.register_buffer("classifier_control_lut", marker_lut, persistent=False) - self.register_buffer("classifier_slot_lut", slot_lut, persistent=False) - else: - self.classifier_control_lut = None - self.classifier_slot_lut = None def split_indices( - self, - adapter_indices: torch.Tensor, - classifier_probe_indices: torch.Tensor | None = None, + self, adapter_indices: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Split ``[total_tokens]`` into (lora_indices, classifier_indices). - See :func:`granite_switch.token_exchange.split_adapter_indices`. + See :func:`granite_switch.classifier_routing.split_adapter_indices`. """ - return split_adapter_indices( - self.adapter_kind_lut, - adapter_indices, - classifier_probe_indices, - ) - - def classifier_indices_from_tokens( - self, input_ids: torch.Tensor - ) -> torch.Tensor | None: - if self.classifier_slot_lut is None: - return None - return self.classifier_slot_lut[input_ids] + return split_adapter_indices(self.adapter_kind_lut, adapter_indices) @property def num_cache_layers(self) -> int: @@ -429,16 +410,10 @@ def forward( # Vectorized control-token matching against ORIGINAL input_ids. matches = input_ids.unsqueeze(1) == adapter_token_ids.unsqueeze(0) # [T, A] is_control_token = matches.any(dim=1) # [total_tokens] - route_control_token = is_control_token - if self.classifier_control_lut is not None: - # A classifier marker is an orthogonal probe. Do not let it replace - # the active adapter stored in the switch's KV-backed coded memory. - is_classifier_marker = self.classifier_control_lut[input_ids] - route_control_token = is_control_token & ~is_classifier_marker # expert_id = argmax + offset (offset selects the layout; see __init__). # 0 for non-control tokens. expert_ids = torch.where( - route_control_token, + is_control_token, matches.long().argmax(dim=1) + self._expert_id_offset, torch.zeros_like(input_ids, dtype=torch.long), ) # [total_tokens] @@ -449,16 +424,16 @@ def forward( _zero = torch.tensor(0.0, dtype=dtype, device=device) _one = torch.tensor(1.0, dtype=dtype, device=device) - # Keys: default masked (-1e9); un-mask (0) at the anchor and adapter - # selection controls. Classifier markers do not advance the persistent - # adapter address. + # Keys: default masked (-1e9); un-mask (0) at anchor + control tokens + # so the one-hot query attends only to those. Arithmetic (torch.where), + # no boolean-index assignment. k_count = torch.full( (total_tokens, self.num_kv_heads, self.counting_head_dim), _NEG_INF, device=device, dtype=dtype, ) - anchor_or_control = is_counting_anchor | route_control_token + anchor_or_control = is_counting_anchor | is_control_token k_count[:, 0, 0] = torch.where(anchor_or_control, _zero, k_count[:, 0, 0]) # Values: v=1 at the position-0 anchor only (the 1/(1+n) numerator). @@ -502,7 +477,7 @@ def forward( # ================================================================== # Look up the codeword for each token's address (compile-safe gather). all_code_vectors = self.codebook[write_addresses].to(dtype) # [T, memory_dim] - write_mask = route_control_token.unsqueeze(-1).to(dtype) # [T, 1] + write_mask = is_control_token.unsqueeze(-1).to(dtype) # [T, 1] # Keys: code(n) * memory_gain at control tokens, zero elsewhere # (arithmetic masking). Value: expert_id at control tokens. @@ -520,7 +495,7 @@ def forward( device=device, dtype=dtype, ) - v_memory[:, 0, 0] = expert_ids.to(dtype) * route_control_token.to(dtype) + v_memory[:, 0, 0] = expert_ids.to(dtype) * is_control_token.to(dtype) # Query: code(n) for every token (reads back the value at address n). q_memory = torch.zeros( diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_classifier_compose.py index bd619022..c4136bc9 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_classifier_compose.py @@ -229,7 +229,7 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): - """[lora, classifier]: retain LoRA at the classifier marker and score it.""" + """[lora, classifier]: LoRA slot 0 populated, classifier slot 1 loaded.""" tokenizer = _tokenizer(MIXED_BASE) hidden = _hidden_size(MIXED_BASE) lora_dir = _find_rag_lora_dir() @@ -292,19 +292,26 @@ def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): "classifier slot 1 LoRA row is not zero — slot numbering leaked" ) - # Select LoRA slot 1, then probe at the classifier's final marker. The - # classifier marker must not replace the active LoRA in the decoder stream. - prompt = _classifier_prompt_ids(tokenizer, clf_ctrl) - input_ids = torch.cat( - [prompt[:, :-1], torch.tensor([[lora_ctrl]]), prompt[:, -1:]], dim=1 - ).to(next(model.parameters()).device) + # A detect forward on the classifier control token emits the winning label. + device = next(model.parameters()).device + input_ids = _classifier_prompt_ids(tokenizer, clf_ctrl).to(device) with torch.no_grad(): out = model(input_ids=input_ids, use_cache=False) emitted = int(out.logits[0, -1].argmax().item()) assert emitted == label_ids[winner] - assert model.model._last_lora_indices[0, -1].item() == 1 + assert model.model._last_lora_indices[0, -1].item() == 0 assert model.model._last_classifier_indices[0, -1].item() == 2 + # The LoRA slot is selected on its own and leaves the LM logits intact. + lora_ids = torch.cat( + [torch.tensor([[lora_ctrl]]), input_ids[:, :-1].cpu()], dim=1 + ).to(device) + with torch.no_grad(): + out = model(input_ids=lora_ids, use_cache=False) + assert model.model._last_lora_indices[0, -1].item() == 1 + assert model.model._last_classifier_verdict is None + assert torch.isfinite(out.logits[0, -1]).all() + def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): """Two classifier slots with DIFFERENT label counts (2 and 3). diff --git a/tests/composer/test_control_lut_refresh.py b/tests/composer/test_control_lut_refresh.py index 9bee8f61..d722969c 100644 --- a/tests/composer/test_control_lut_refresh.py +++ b/tests/composer/test_control_lut_refresh.py @@ -43,7 +43,7 @@ _VOCAB_WITH_AUDIO = 203 -def _config(switch_type): +def _config(switch_type, **overrides): return GraniteSwitchConfig( vocab_size=_BASE_VOCAB, hidden_size=32, @@ -63,6 +63,7 @@ def _config(switch_type): adapter_substitute_token_ids=list(_SUB_IDS), switch_type=switch_type, torch_dtype=torch.float32, + **overrides, ) @@ -129,3 +130,41 @@ def test_no_mapping_is_not_an_error(switch_type): assert model.model.switch.control_to_substitute_lut is None assert refresh_switch_control_lut(model) is False + + +def test_classifier_routing_survives_audio_resize_and_strict_reload(tmp_path): + """Classifier + audio: every retained routing table must match the saved config. + + The control LUT is vocabulary-sized and is rebuilt after the resize; the + classifier kind LUT is indexed by adapter, so the resize cannot stale it. + A strict reload (no ``ignore_mismatched_sizes``) must keep both, and the + classifier marker must still route to its head and emit a label word. + """ + from granite_switch.composer.validator import validate_control_lut + + config = _config( + "multi", + adapter_kinds=["lora", "classifier"], + classifier_label_token_ids=[None, [10, 11]], + ) + model = GraniteSwitchForCausalLM(config).eval() + with torch.no_grad(): + model.model.classifier_head.bias[1, :2] = torch.tensor([-5.0, 5.0]) + model.resize_token_embeddings(_VOCAB_WITH_AUDIO) + assert refresh_switch_control_lut(model) is True + validate_control_lut(model) + + kind_lut = model.model.switch.adapter_kind_lut + assert kind_lut.tolist() == [False, False, True] + + model.save_pretrained(tmp_path / "composed") + loaded = GraniteSwitchForCausalLM.from_pretrained(tmp_path / "composed").eval() + switch = loaded.model.switch + assert switch.control_to_substitute_lut.numel() == loaded.config.vocab_size + torch.testing.assert_close(switch.adapter_kind_lut, kind_lut) + + with torch.no_grad(): + out = loaded(input_ids=torch.tensor([[1, 2, 3, _CTRL_IDS[1]]]), use_cache=False) + assert loaded.model._last_lora_indices.tolist() == [[0, 0, 0, 0]] + assert loaded.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] + assert int(out.logits[0, -1].argmax()) == 11 diff --git a/tests/composer/test_sr_compose_e2e.py b/tests/composer/test_sr_compose_e2e.py index 90c542f8..d9f9c316 100644 --- a/tests/composer/test_sr_compose_e2e.py +++ b/tests/composer/test_sr_compose_e2e.py @@ -187,8 +187,8 @@ def _create_sr_adapter( def _compose_sr_classifier(path: Path, read_layers=(None,)): """Tiny SR + classifier checkpoint, also used by the vLLM test. - An identity head exposes every component of its input, making a base-stream - read distinguishable from an adapted-stream read without random margins. + An identity head exposes every component of its input, so its scores pin + the exact layer and normalization read without random margins. """ from granite_switch.composer import GraniteSwitchComposer from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE @@ -221,8 +221,12 @@ def _compose_sr_classifier(path: Path, read_layers=(None,)): @pytest.mark.parametrize("read_layer", [None, 0, NUM_LAYERS - 1]) -def test_sr_classifier_reads_selected_stream(tmp_path, capsys, read_layer): - """Identity scores prove stream, layer and normalization before/after reload.""" +def test_sr_and_classifier_slots_select_separately(tmp_path, capsys, read_layer): + """Identity scores prove layer and normalization before/after reload. + + The classifier marker routes out of the SR stream, so the probe reads the + base stream at its layer; an SR request in the same checkpoint still adapts. + """ from granite_switch.hf import load_model model = _compose_sr_classifier(tmp_path, read_layers=[read_layer]) @@ -235,62 +239,50 @@ def test_sr_classifier_reads_selected_stream(tmp_path, capsys, read_layer): assert torch.count_nonzero(layer.cross_stream.lora_B[0]) > 0 assert torch.count_nonzero(layer.cross_stream.lora_B[1]) == 0 - streams = [] - hook = model.model.layers[ - -1 if read_layer is None else read_layer - ].register_forward_hook(lambda _module, _args, output: streams.append(output[0])) - ids = torch.tensor([[10, CONTROL_TOKEN_ID, 20, CONTROL_TOKEN_ID + 1]]) - try: - with torch.no_grad(): - output = model(input_ids=ids, use_cache=False, output_hidden_states=True) - finally: - hook.remove() - assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1]] - assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] - base, adapted = ( - (model.model.norm(h) if read_layer is None else h)[0, -1] for h in streams[0] - ) - assert not torch.allclose(adapted, base), "fixture must exercise SR adaptation" - if read_layer is None: - torch.testing.assert_close(output.hidden_states[-1][0, -1], adapted) - else: - assert not torch.allclose(adapted, output.hidden_states[-1][0, -1]) + def run(m, ids): + streams = [] + hook = m.model.layers[ + -1 if read_layer is None else read_layer + ].register_forward_hook(lambda _mod, _args, out: streams.append(out[0])) + try: + with torch.no_grad(): + output = m(input_ids=ids, use_cache=False) + finally: + hook.remove() + base, adapted = ( + (m.model.norm(h) if read_layer is None else h)[0, -1] for h in streams[0] + ) + return output, base, adapted + labels = model.config.classifier_label_token_ids[1] - torch.testing.assert_close(output.logits[0, -1, labels], adapted) non_labels = torch.ones(VOCAB_SIZE, dtype=torch.bool) non_labels[labels] = False - assert torch.isneginf(output.logits[0, -1, non_labels]).all() + + probe_ids = torch.tensor([[10, 21, 20, CONTROL_TOKEN_ID + 1]]) + probe, base, _ = run(model, probe_ids) + assert model.model._last_lora_indices.tolist() == [[0, 0, 0, 0]] + assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] + torch.testing.assert_close(probe.logits[0, -1, labels], base) + assert torch.isneginf(probe.logits[0, -1, non_labels]).all() + + sr_ids = torch.tensor([[10, CONTROL_TOKEN_ID, 20, 30]]) + sr, base, adapted = run(model, sr_ids) + assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1]] + assert not model.model._last_classifier_indices.any() + assert not torch.allclose(adapted, base), "fixture must exercise SR adaptation" + assert torch.isfinite(sr.logits[0, -1]).all() checkpoint = tmp_path / "composed" model.save_pretrained(checkpoint) loaded = load_model(str(checkpoint)).eval() - for name in ("adapter_kind_lut", "classifier_control_lut", "classifier_slot_lut"): - torch.testing.assert_close( - getattr(loaded.model.switch, name), getattr(model.model.switch, name) - ) - with torch.no_grad(): - restored = loaded(input_ids=ids, use_cache=False) - torch.testing.assert_close(restored.logits, output.logits) - assert loaded.model._last_lora_indices[0, -1].item() == 1 - assert loaded.model._last_classifier_indices[0, -1].item() == 2 - - # With no SR activation, the probe must read the base stream at that layer. - streams.clear() - hook = model.model.layers[ - -1 if read_layer is None else read_layer - ].register_forward_hook(lambda _module, _args, result: streams.append(result[0])) - try: + torch.testing.assert_close( + loaded.model.switch.adapter_kind_lut, model.model.switch.adapter_kind_lut + ) + for ids, expected in ((probe_ids, probe), (sr_ids, sr)): with torch.no_grad(): - inactive = model( - input_ids=torch.tensor([[10, 21, 20, CONTROL_TOKEN_ID + 1]]), - use_cache=False, - ) - finally: - hook.remove() - expected = streams[0][0] - if read_layer is None: - expected = model.model.norm(expected) - torch.testing.assert_close(inactive.logits[0, -1, labels], expected[0, -1]) + restored = loaded(input_ids=ids, use_cache=False) + torch.testing.assert_close(restored.logits, expected.logits) + assert loaded.model._last_lora_indices[0, -1].item() == 1 # --------------------------------------------------------------------------- diff --git a/tests/hf/test_classifier.py b/tests/hf/test_classifier.py index e94baacf..e0fa9df1 100644 --- a/tests/hf/test_classifier.py +++ b/tests/hf/test_classifier.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Classifier-substitute path tests for GraniteSwitchForCausalLM (HF backend). -A classifier slot is selected independently from the active adapter. A LoRA -selected earlier in the prompt remains active at the classifier marker, while -the classifier head reads its slot's read layer (the final hidden state by -default) to emit ONE verdict per request. The verdict marker must be the -request's last token. +A classifier slot is an alternative to a LoRA adapter: it shares the adapter +index/control-token machinery, the switch splits classifier positions out of +the LoRA stream, and the classifier head reads its slot's read layer (the final +hidden state by default) to emit ONE verdict per request. The verdict is read at +the control token, which must be the request's last token. The verdict exits as a GENERATED LABEL WORD, matching the vLLM backend: the LM head rewrites the marker's logit row so the vocab is -inf except at @@ -97,19 +97,25 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): assert int(last_row.argmax()) in label_ids -def test_classifier_marker_preserves_active_lora(): +def test_lora_and_classifier_slots_select_separately(): config = _classifier_config(["lora", "classifier"]) model = GraniteSwitchForCausalLM(config).eval() _set_adapter_token_ids(model, config.adapter_token_ids) - # Select LoRA slot 1, then place the classifier marker last. The marker - # probes slot 2 while LoRA slot 1 remains active through the final token. - input_ids = torch.tensor([[10, 250, 20, 30, 251]]) + # The LoRA control token selects slot 1 for the following tokens. + with torch.no_grad(): + model(input_ids=torch.tensor([[10, 250, 20, 30]])) + assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1]] + assert not model.model._last_classifier_indices.any() + assert model.model._last_classifier_verdict is None + + # The classifier marker selects slot 2 and routes out of the LoRA stream. + input_ids = torch.tensor([[10, 20, 30, 251]]) with torch.no_grad(): output = model(input_ids=input_ids) - assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1, 1]] - assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 0, 2]] + assert model.model._last_lora_indices.tolist() == [[0, 0, 0, 0]] + assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] label_ids = config.classifier_label_token_ids[1] last_logits = output.logits[0, -1] assert torch.isfinite(last_logits[torch.tensor(label_ids)]).all() @@ -325,13 +331,11 @@ def test_intermediate_layer_classifier_scores_that_layers_output(): _randomize_classifier_head(model) with torch.no_grad(): out = model( - input_ids=torch.tensor([[10, 250, 30, 40, 50, 60, 70, 251]]), + input_ids=torch.tensor([[10, 20, 30, 40, 50, 60, 70, 251]]), output_hidden_states=True, ) - # The classifier marker selects its head independently; the active LoRA - # remains routed through the layer whose residual representation is read. - assert model.model._last_lora_indices[0, -1].item() == 1 + assert model.model._last_lora_indices[0, -1].item() == 0 assert model.model._last_classifier_indices[0, -1].item() == 2 rows, pos, verdict, slots, _ = model.model._last_classifier_verdict expected = model.model.classifier_head(out.hidden_states[1][rows, pos], slots) diff --git a/tests/unit/test_classifier_routing.py b/tests/unit/test_classifier_routing.py new file mode 100644 index 00000000..bdee44ed --- /dev/null +++ b/tests/unit/test_classifier_routing.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Classifier-slot routing: the kind LUT and the LoRA/classifier index split.""" + +import torch + +from granite_switch.classifier_routing import ( + build_adapter_kind_lut, + split_adapter_indices, +) +from granite_switch.config import GraniteSwitchConfig + + +def _config(adapter_kinds, label_ids, base_reset=False): + num_adapters = len(adapter_kinds) + token_ids = list(range(500, 500 + num_adapters)) + if base_reset: + token_ids = [499, *token_ids] + return GraniteSwitchConfig( + vocab_size=600, + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=4, + num_hidden_layers=4, + intermediate_size=128, + num_adapters=num_adapters, + adapter_token_ids=token_ids, + adapter_names=[f"a{i}" for i in range(num_adapters)], + adapter_substitute_token_ids=[1] * len(token_ids), + adapter_ranks=[4] * num_adapters, + max_lora_rank=4, + adapter_kinds=adapter_kinds, + classifier_label_token_ids=label_ids, + ) + + +class TestAdapterKindLut: + def test_none_without_classifier_slots(self): + assert build_adapter_kind_lut(_config(["lora", "lora"], None)) is None + + def test_marks_classifier_adapter_indices(self): + cfg = _config(["lora", "classifier", "classifier"], [None, [100], [101, 102]]) + assert build_adapter_kind_lut(cfg).tolist() == [False, False, True, True] + + def test_base_reset_token_does_not_shift_slots(self): + """The LUT is indexed by adapter index, not by control-token position.""" + cfg = _config(["lora", "classifier"], [None, [100, 101]], base_reset=True) + assert build_adapter_kind_lut(cfg).tolist() == [False, False, True] + assert cfg.classifier_control_token_ids == [501] + + +class TestSplitAdapterIndices: + def test_classifier_positions_leave_the_lora_stream(self): + cfg = _config(["lora", "classifier"], [None, [100, 101]]) + lora, classifier = split_adapter_indices( + build_adapter_kind_lut(cfg), torch.tensor([0, 1, 1, 0, 2]) + ) + assert lora.tolist() == [0, 1, 1, 0, 0] + assert classifier.tolist() == [0, 0, 0, 0, 2] + + def test_preserves_batched_shape(self): + cfg = _config(["classifier", "lora"], [[100, 101], None]) + lora, classifier = split_adapter_indices( + build_adapter_kind_lut(cfg), torch.tensor([[0, 2, 1], [1, 1, 0]]) + ) + assert lora.tolist() == [[0, 2, 0], [0, 0, 0]] + assert classifier.tolist() == [[0, 0, 1], [1, 1, 0]] + + def test_no_classifier_slots_pass_through(self): + indices = torch.tensor([0, 1, 2]) + lora, classifier = split_adapter_indices(None, indices) + assert lora is indices + assert not classifier.any() diff --git a/tests/unit/test_token_exchange.py b/tests/unit/test_token_exchange.py index bde219f9..f008046c 100644 --- a/tests/unit/test_token_exchange.py +++ b/tests/unit/test_token_exchange.py @@ -6,14 +6,8 @@ """ import pytest -import torch from granite_switch.config import GraniteSwitchConfig -from granite_switch.token_exchange import ( - build_adapter_kind_lut, - build_classifier_control_luts, - split_adapter_indices, -) def _base(num_adapters=2, **overrides): @@ -66,40 +60,3 @@ class TestProjectionHeadDim: def test_inferred_from_hidden_size(self): cfg = GraniteSwitchConfig(**_base()) assert cfg.projection_head_dim == cfg.hidden_size // cfg.num_attention_heads - - -class TestClassifierLoRACoactivation: - def _config(self): - return GraniteSwitchConfig( - **_base( - adapter_token_ids=[499, 500, 501], # leading base-reset token - adapter_substitute_token_ids=[1, 1, 1], - adapter_kinds=["lora", "classifier"], - classifier_label_token_ids=[None, [100, 101]], - ) - ) - - def test_classifier_probe_lut_skips_leading_base_reset_token(self): - cfg = self._config() - marker_lut, slot_lut = build_classifier_control_luts(cfg) - - assert not marker_lut[499] - assert not marker_lut[500] - assert marker_lut[501] - assert slot_lut[499] == slot_lut[500] == 0 - assert slot_lut[501] == 2 - assert cfg.classifier_control_token_ids == [501] - - def test_classifier_probe_preserves_lora_indices_and_separates_probe(self): - cfg = self._config() - adapter_indices = torch.tensor([0, 1, 1, 1]) - probe_indices = torch.tensor([0, 0, 0, 2]) - - lora_indices, classifier_indices = split_adapter_indices( - build_adapter_kind_lut(cfg), - adapter_indices, - probe_indices, - ) - - assert lora_indices.tolist() == [0, 1, 1, 1] - assert classifier_indices.tolist() == [0, 0, 0, 2] diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py index d6c8fd45..d60f8953 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -399,7 +399,11 @@ def run_at_layer(ckpt, prompt_ids, label_token_ids, out_path, full_cudagraph=Fal def run_sr_classifier(ckpt, out_path, *, eager, pp_size=1): - """Exercise SR probes at each configured layer, sequentially and in a mixed batch.""" + """Serve each SR-checkpoint probe and an SR request, sequentially and mixed. + + Each slot is selected on its own: a probe request carries only its marker, + and the SR request carries only the SR control token. + """ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" if pp_size > 1: # Only trusted, importable test callbacks cross the local RPC boundary. @@ -437,7 +441,7 @@ def run_sr_classifier(ckpt, out_path, *, eager, pp_size=1): llm = LLM(**kw) _record_classifier_rows(llm) sampling = SamplingParams(max_tokens=1, temperature=0.0, logprobs=HIDDEN + 1) - plain = [10, CONTROL_TOKEN_ID, 20, 30] + sr_request = [10, CONTROL_TOKEN_ID, 20, 30] def serve(prompts): outputs = llm.generate( @@ -463,20 +467,15 @@ def serve(prompts): ], } + # Each probe twice: the second run is the replay of the first's shape. result = { "sequential": [ - serve([p]) - for marker in markers - for p in ( - [10, CONTROL_TOKEN_ID, 20, marker], - [10, 21, 20, marker], - [10, CONTROL_TOKEN_ID, 20, marker], - ) + serve([[10, 21, 20, marker]]) for marker in markers for _ in range(2) ] } - mixed = [[10, CONTROL_TOKEN_ID, 20, markers[0]], plain, [10, 21, 20, markers[-1]]] + mixed = [[10, 21, 20, markers[0]], sr_request, [11, 21, 20, markers[-1]]] if len(markers) > 1: - mixed.append([10, CONTROL_TOKEN_ID, 20, markers[1]]) + mixed.append([12, 21, 20, markers[1]]) result["mixed"] = serve(mixed) with open(out_path, "w") as f: json.dump(result, f) diff --git a/tests/vllm/_model_forward_tests.py b/tests/vllm/_model_forward_tests.py index 1ec442dc..afc4ac1a 100644 --- a/tests/vllm/_model_forward_tests.py +++ b/tests/vllm/_model_forward_tests.py @@ -474,23 +474,21 @@ def test_different_adapters_produce_different_logits(self): ) -class TestClassifierLoRAProbe(_VLLMModelTestBase): +class TestClassifierSlotRouting(_VLLMModelTestBase): def _make_config(self): return _tiny_vllm_config( adapter_kinds=["lora", "classifier"], classifier_label_token_ids=[None, [100, 101]], ) - def test_marker_scores_with_active_lora_and_emits_label(self, monkeypatch): + def test_marker_routes_out_of_lora_and_emits_label(self, monkeypatch): self.model.eval() seen = {} switch = self.model.model.switch original_split = switch.split_indices - def record_split(adapter_indices, classifier_probe_indices=None): - lora_indices, classifier_indices = original_split( - adapter_indices, classifier_probe_indices - ) + def record_split(adapter_indices): + lora_indices, classifier_indices = original_split(adapter_indices) seen["lora"] = lora_indices.detach().cpu() seen["classifier"] = classifier_indices.detach().cpu() return lora_indices, classifier_indices @@ -498,10 +496,10 @@ def record_split(adapter_indices, classifier_probe_indices=None): monkeypatch.setattr(switch, "split_indices", record_split) with torch.no_grad(): - logits = self._run_forward_and_logits([10, 250, 20, 30, 251]) + logits = self._run_forward_and_logits([10, 20, 30, 251]) - assert seen["lora"].tolist() == [0, 1, 1, 1, 1] - assert seen["classifier"].tolist() == [0, 0, 0, 0, 2] + assert seen["lora"].tolist() == [0, 0, 0, 0] + assert seen["classifier"].tolist() == [0, 0, 0, 2] final_logits = logits[-1] assert torch.isfinite(final_logits[[100, 101]]).all() non_label_ids = [ diff --git a/tests/vllm/test_classifier_capture_safety.py b/tests/vllm/test_classifier_capture_safety.py index 00bd2ae2..31f1c32b 100644 --- a/tests/vllm/test_classifier_capture_safety.py +++ b/tests/vllm/test_classifier_capture_safety.py @@ -205,21 +205,29 @@ def test_pipeline_handoff_without_later_token_ids(monkeypatch, sr, stages, dtype A layer-0 probe crosses the boundary; another probe and final norm are scored on the last rank. With three stages, the middle rank only relays scores. Distinct values expose resets, wrong layers, stream selection and reliance - on input_ids. The token slice matches vLLM's PP buffer contract. + on input_ids. The switch keeps the latest slot, like MultiSwitch, so only + the markers themselves may carry a classifier index. Markers route out of + the LoRA stream; the trailing LoRA token reads the adapted SR stream. The + token slice matches vLLM's PP buffer contract. """ from granite_switch.vllm import granite_switch_model as model_module class Switch: def __call__(self, input_ids, **kwargs): - return (input_ids != 5).long(), input_ids - - def classifier_indices_from_tokens(self, input_ids): - lut = torch.zeros(100, dtype=torch.long) - lut[90:93] = torch.tensor([2, 3, 4]) - return lut[input_ids] - - def split_indices(self, indices, probes): - return indices, probes + control_slot = {89: 1, 90: 2, 91: 3, 92: 4} + active, indices = 0, [] + for token in input_ids.tolist(): + active = control_slot.get(token, active) + indices.append(active) + return torch.tensor(indices), input_ids + + def split_indices(self, indices): + is_classifier = indices >= 2 + zeros = torch.zeros_like(indices) + return ( + torch.where(is_classifier, zeros, indices), + torch.where(is_classifier, indices, zeros), + ) class Layer(torch.nn.Module): def forward(self, positions, hidden_states, residual=None): @@ -240,6 +248,7 @@ def forward(self, positions, hidden_states, residual=None): model.config.fused_add_norm = False model.switch = Switch() model.adapter_token_ids = torch.tensor([89, 90, 91, 92]) + model.classifier_control_token_ids_t = torch.tensor([90, 91, 92]) model.lora_meta = model.lora_ctx = None model.decoder_interface = SRDecoderInterface() if sr else LoRADecoderInterface() model.layers = torch.nn.ModuleList([Layer() for _ in range(stages)]) @@ -255,13 +264,13 @@ def forward(self, positions, hidden_states, residual=None): ) wire = GraniteSwitchModel.forward( model, - torch.tensor([5, 90, 91, 92, 5]), + torch.tensor([5, 90, 91, 92, 89]), torch.arange(5), inputs_embeds=torch.ones(5, 4, dtype=dtype), )[1:] assert isinstance(wire, IntermediateTensors) assert wire["classifier_indices"].tolist() == [2, 3, 4, 0] - assert wire["classifier_scores"][:, 0].tolist() == [22 if sr else 4, 0, 0, 0] + assert wire["classifier_scores"][:, 0].tolist() == [4, 0, 0, 0] buffers = GraniteSwitchModel.make_empty_intermediate_tensors( model, 4, dtype, torch.device("cpu") ) @@ -283,10 +292,10 @@ def forward(self, positions, hidden_states, residual=None): for key, tensor in packed.items(): buffers[key].copy_(tensor) assert packed[:, -1].tolist() == [2, 3, 4, 0] - source = 1 + (10 if sr else 1) * stages - expected = [22 if sr else 4, 3 * source, 4 * source, 0] - assert packed[:, 4].tolist() == expected - assert packed[-1, 0] == stages + 1 # inactive SR token reads the base stream + # Markers read the base stream; the LoRA token reads the adapted SR stream. + source = 1 + stages + assert packed[:, 4].tolist() == [4, 3 * source, 4 * source, 0] + assert packed[-1, 0] == 1 + (10 if sr else 1) * stages def test_sr_pipeline_stack_partitions_post_attention_metadata(): diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_classifier_chunked_prefill.py index aa70b5f9..09fb2e62 100644 --- a/tests/vllm/test_classifier_chunked_prefill.py +++ b/tests/vllm/test_classifier_chunked_prefill.py @@ -131,18 +131,17 @@ def composed_sr_probes(tmp_path_factory): @pytest.mark.parametrize("mode", ["eager", "full"]) -def test_sr_and_final_classifier_coactivate(composed_sr, mode): - """The identity head scores the selected SR output; transport stays exact.""" +def test_sr_checkpoint_serves_final_classifier_slot(composed_sr, mode): + """The identity head scores the served hidden state; transport stays exact.""" served = _run_worker( composed_sr["outdir"], [f"run-sr-{mode}"], {"CLS_CKPT": composed_sr["ckpt"]}, ) sequential = served["sequential"] - assert [len(case["records"]) for case in sequential] == [1, 1, 1] - active, inactive, replay = [case["records"][0] for case in sequential] - assert active["served_hidden"] != inactive["served_hidden"] - assert replay["packed_scores"] == active["packed_scores"] + assert [len(case["records"]) for case in sequential] == [1, 1] + probe, replay = [case["records"][0] for case in sequential] + assert replay["packed_scores"] == probe["packed_scores"] for case in [*sequential, served["mixed"]]: for record in case["records"]: assert record["slot"] == 2 @@ -175,23 +174,18 @@ def sr_probe_reference(composed_sr_probes, request): def _check_sr_layer_probes(result, mode, pp): - """Distinct identity-head rows expose wrong sources, masks and stale replay.""" + """Distinct identity-head rows expose wrong layers, masks and stale replay.""" failures = [] - assert len(result["sequential"]) == 9 - active_scores = [] + assert len(result["sequential"]) == 6 + slot_scores = [] for slot in (2, 3, 4): - active, inactive, replay = result["sequential"][3 * (slot - 2) : 3 * (slot - 1)] - assert [r["slot"] for r in active["records"]] == [slot] - assert [r["slot"] for r in inactive["records"]] == [slot] - scores = active["records"][0]["packed_scores"] - active_scores.append(scores) - assert scores != inactive["records"][0]["packed_scores"] + probe, replay = result["sequential"][2 * (slot - 2) : 2 * (slot - 1)] + assert [r["slot"] for r in probe["records"]] == [slot] + assert [r["slot"] for r in replay["records"]] == [slot] + scores = probe["records"][0]["packed_scores"] + slot_scores.append(scores) assert scores == replay["records"][0]["packed_scores"] - for name, case in ( - ("active", active), - ("inactive", inactive), - ("replay", replay), - ): + for name, case in (("probe", probe), ("replay", replay)): record, output = case["records"][0], case["outputs"][0] failures.extend( _check_transport( @@ -204,9 +198,7 @@ def _check_sr_layer_probes(result, mode, pp): max(record["packed_scores"]) ) assert not output["has_non_label_support"] - assert all( - active_scores[a] != active_scores[b] for a, b in ((0, 1), (0, 2), (1, 2)) - ) + assert all(slot_scores[a] != slot_scores[b] for a, b in ((0, 1), (0, 2), (1, 2))) assert sorted(r["slot"] for r in result["mixed"]["records"]) == [2, 3, 4] assert [o["has_non_label_support"] for o in result["mixed"]["outputs"]] == [ False, @@ -218,7 +210,7 @@ def _check_sr_layer_probes(result, mode, pp): def test_sr_classifier_reads_each_layer_under_replay(sr_probe_reference): - """One GPU: probes at layer 0, layer 1 and final norm with SR active/inactive.""" + """One GPU: probes at layer 0, layer 1 and final norm beside an SR slot.""" mode, served = sr_probe_reference _check_sr_layer_probes(served, mode, 1) diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index 09979e8b..fed620d1 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -9,8 +9,12 @@ serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detecti A Granite Switch adapter slot is normally a LoRA adapter. A **classifier slot** is an alternative occupant of that same slot: instead of a LoRA, it holds a small linear classification head. It is triggered by the *same* control-token machinery — one `<|name|>` token, one adapter index — but -the classifier reads the configured layer's activation and produces per-label logits. An active -LoRA or SR adapter stays active while the classifier reads its resulting hidden state. +the classifier reads the configured layer's activation and produces per-label logits. + +Activation probes are supported as classifier slots in the standard composition workflow. Probes and +generation adapters can coexist in one checkpoint and be selected individually: a request selects +one slot, either a generation adapter or a classifier. The calling stack manages compatible +execution state and cache reuse. At inference the verdict does not come back as a raw score vector. Each label word is resolved to a single token id at compose time, and the model rewrites the request's final logit row so that the @@ -145,10 +149,9 @@ safety_mid: layer: 20 # 0..num_decoder_layers - 1; omit for final ``` -This also works with Shadow Residual (SR) adapters: an active adapter's tokens read the adapted -stream after that layer's cross-stream shunt; inactive tokens read the base stream. Decoder layer -numbers exclude the switch's KV-cache slots. Pipeline parallelism preserves the same numbering -across stages. +Classifier slots can also be composed alongside Shadow Residual (SR) adapters; a classifier +request reads that layer's base stream. Decoder layer numbers exclude the switch's KV-cache slots. +Pipeline parallelism preserves the same numbering across stages. You can mix LoRA and classifier entries in one manifest, and give each classifier its own `labels` list of any length: @@ -331,8 +334,9 @@ channel. ## How it works (brief) -- **Independent probe trigger.** One `<|name|>` control token selects a classifier head while - preserving the active LoRA or SR adapter. The head reads the resulting activation at that token. +- **Same trigger, separate path.** One `<|name|>` control token selects a classifier slot, exactly + like an adapter. Classifier positions are routed out of the LoRA stream, and the head reads the + activation at that token. - **Padded head bank.** All classifier heads live in one stacked tensor padded to `max_classifier_labels`; a slot with fewer labels leaves its extra rows zero. This keeps the vLLM head a single fused, compile-safe operation regardless of how many labels each slot has — the same From 4c32bf8ebc7ac0d0f940fbc3c5ae598bd9246495 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 21:26:50 +0000 Subject: [PATCH 15/20] test(classifier): compose a classifier slot with audio and reload strictly Real compose of a classifier manifest with and without --enable-audio, then a strict HF load that checks the routing tables and a deterministic verdict. Covers the reported classifier + audio load failure end to end. Signed-off-by: IRENEKO --- .../composer/test_multi_audio_compose_e2e.py | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/tests/composer/test_multi_audio_compose_e2e.py b/tests/composer/test_multi_audio_compose_e2e.py index 1b4d24e3..a3ba27f8 100644 --- a/tests/composer/test_multi_audio_compose_e2e.py +++ b/tests/composer/test_multi_audio_compose_e2e.py @@ -153,3 +153,121 @@ def test_multi_audio_checkpoint_loads_without_size_overrides(audio_multi): assert lut[ctrl_id].item() == sub_id, ctrl_id # Every non-control id stays a sentinel, so ordinary tokens pass through. assert int((lut >= 0).sum()) == len(model.config.adapter_token_ids) + + +# --------------------------------------------------------------------------- # +# Classifier slot + audio: every vocabulary-sized routing table must be rebuilt +# after the audio marker grows the vocabulary, and the classifier must still +# route to its head after a strict reload. Text-only compose is the control. +# --------------------------------------------------------------------------- # +_CLASSIFIER_BASE = "ibm-granite/granite-4.1-3b" +_CLASSIFIER_LABELS = ["safe", "unsafe"] + + +def _write_classifier_manifest(root): + """A synthetic bias-only head (label 1 always wins) under the library layout.""" + import torch + from safetensors.torch import save_file + from transformers import AutoConfig + + hidden = AutoConfig.from_pretrained(_CLASSIFIER_BASE).hidden_size + head_dir = root / "safety" / "safety" / "granite-4.1-3b" / "lora" + head_dir.mkdir(parents=True) + bias = torch.tensor([-10.0, 10.0]) + save_file( + {"weight": torch.zeros(len(_CLASSIFIER_LABELS), hidden), "bias": bias}, + str(head_dir / "classifier_head.safetensors"), + ) + (head_dir / "io.yaml").write_text("name: safety\nmodel: ~\n") + manifest = root / "adapters.yaml" + manifest.write_text( + "safety:\n" + f" path: {head_dir}\n" + " type: lora\n" + " kind: classifier\n" + f" labels: [{', '.join(_CLASSIFIER_LABELS)}]\n" + ) + return manifest + + +def _compose_classifier(root, enable_audio): + manifest = _write_classifier_manifest(root) + output_dir = root / "model" + cmd = [ + sys.executable, + "-m", + "granite_switch.composer.compose_granite_switch", + "--base-model", + _CLASSIFIER_BASE, + "--adapters", + str(manifest), + "--output", + str(output_dir), + ] + if enable_audio: + cmd.append("--enable-audio") + result = subprocess.run(cmd, capture_output=True, text=True, timeout=BUILD_TIMEOUT) + assert result.returncode == 0, ( + f"classifier compose failed (exit {result.returncode}).\n" + f"STDOUT tail:\n{result.stdout[-2000:]}\n" + f"STDERR tail:\n{result.stderr[-2000:]}" + ) + return output_dir + + +@pytest.fixture(scope="module", params=[True, False], ids=["audio", "text"]) +def classifier_checkpoint(request, tmp_path_factory): + root = tmp_path_factory.mktemp( + "classifier-audio" if request.param else "classifier" + ) + return request.param, _compose_classifier(root, enable_audio=request.param) + + +@pytest.mark.xdist_group("multi_audio_compose_e2e") +def test_classifier_checkpoint_loads_strictly_and_routes(classifier_checkpoint): + torch = pytest.importorskip("torch") + from transformers import AutoTokenizer + + from granite_switch.hf import GraniteSwitchForCausalLM + + enable_audio, output_dir = classifier_checkpoint + config = json.loads((output_dir / "config.json").read_text()) + if enable_audio: + # Guard the guard: the marker must push the vocabulary past the controls. + assert config["vocab_size"] > max(config["adapter_token_ids"]) + 1 + _assert_lut_matches(output_dir, config["vocab_size"]) + + model = GraniteSwitchForCausalLM.from_pretrained( + str(output_dir), dtype=torch.float32 + ).eval() + switch = model.model.switch + assert switch.control_to_substitute_lut.numel() == model.config.vocab_size + assert switch.adapter_kind_lut.tolist() == [False, True] + + tokenizer = AutoTokenizer.from_pretrained(str(output_dir)) + input_ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "Is this text safe? The weather is nice."}], + adapter_name="safety", + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + )["input_ids"] + marker = tokenizer.convert_tokens_to_ids("<|safety|>") + assert input_ids[0, -1].item() == marker + with torch.no_grad(): + out = model(input_ids=input_ids, use_cache=False) + assert model.model._last_classifier_indices[0, -1].item() == 1 + assert not model.model._last_lora_indices.any() + label_ids = model.config.classifier_label_token_ids[0] + assert int(out.logits[0, -1].argmax()) == label_ids[1] + + +def _assert_lut_matches(output_dir, vocab_size): + from safetensors import safe_open + + for path in sorted(output_dir.glob("*.safetensors")): + with safe_open(str(path), framework="pt") as f: + if _LUT_KEY in f.keys(): + assert f.get_slice(_LUT_KEY).get_shape()[0] == vocab_size + return + pytest.fail(f"{_LUT_KEY} not found in {output_dir}") From 0644cfae480319c6154f6d93b3b40c535fab1141 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 22:13:58 +0000 Subject: [PATCH 16/20] test(classifier): compose the strict-load check as a mixed LoRA + classifier checkpoint Probes ship alongside generation adapters, so the check composes that shape, with and without --enable-audio. Signed-off-by: IRENEKO --- .../composer/test_multi_audio_compose_e2e.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/composer/test_multi_audio_compose_e2e.py b/tests/composer/test_multi_audio_compose_e2e.py index a3ba27f8..339644d9 100644 --- a/tests/composer/test_multi_audio_compose_e2e.py +++ b/tests/composer/test_multi_audio_compose_e2e.py @@ -156,20 +156,23 @@ def test_multi_audio_checkpoint_loads_without_size_overrides(audio_multi): # --------------------------------------------------------------------------- # -# Classifier slot + audio: every vocabulary-sized routing table must be rebuilt -# after the audio marker grows the vocabulary, and the classifier must still -# route to its head after a strict reload. Text-only compose is the control. +# A LoRA + classifier checkpoint composed with and without --enable-audio must +# load strictly (no ignore_mismatched_sizes) and still route the classifier +# marker to its head. The audio marker grows the vocabulary past the control +# ids, so this also checks the saved control LUT matches the shipped config. # --------------------------------------------------------------------------- # _CLASSIFIER_BASE = "ibm-granite/granite-4.1-3b" _CLASSIFIER_LABELS = ["safe", "unsafe"] def _write_classifier_manifest(root): - """A synthetic bias-only head (label 1 always wins) under the library layout.""" + """A real LoRA plus a synthetic bias-only head (label 1 always wins).""" import torch from safetensors.torch import save_file from transformers import AutoConfig + from .test_classifier_compose import _find_rag_lora_dir + hidden = AutoConfig.from_pretrained(_CLASSIFIER_BASE).hidden_size head_dir = root / "safety" / "safety" / "granite-4.1-3b" / "lora" head_dir.mkdir(parents=True) @@ -181,6 +184,9 @@ def _write_classifier_manifest(root): (head_dir / "io.yaml").write_text("name: safety\nmodel: ~\n") manifest = root / "adapters.yaml" manifest.write_text( + "rag:\n" + f" path: {_find_rag_lora_dir()}\n" + " type: lora\n" "safety:\n" f" path: {head_dir}\n" " type: lora\n" @@ -242,7 +248,7 @@ def test_classifier_checkpoint_loads_strictly_and_routes(classifier_checkpoint): ).eval() switch = model.model.switch assert switch.control_to_substitute_lut.numel() == model.config.vocab_size - assert switch.adapter_kind_lut.tolist() == [False, True] + assert switch.adapter_kind_lut.tolist() == [False, False, True] tokenizer = AutoTokenizer.from_pretrained(str(output_dir)) input_ids = tokenizer.apply_chat_template( @@ -256,9 +262,9 @@ def test_classifier_checkpoint_loads_strictly_and_routes(classifier_checkpoint): assert input_ids[0, -1].item() == marker with torch.no_grad(): out = model(input_ids=input_ids, use_cache=False) - assert model.model._last_classifier_indices[0, -1].item() == 1 + assert model.model._last_classifier_indices[0, -1].item() == 2 assert not model.model._last_lora_indices.any() - label_ids = model.config.classifier_label_token_ids[0] + label_ids = model.config.classifier_label_token_ids[1] assert int(out.logits[0, -1].argmax()) == label_ids[1] From a86cbba0303c747c5c28042581f6fc579c70f183 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 22:14:00 +0000 Subject: [PATCH 17/20] refactor(classifier): trim leftover classifier plumbing LoRAContext.classifier_indices has no reader since the vLLM verdict rides the packed hidden states, so lora_kernel_meta.py returns to main. Reuse lora_adapter_names for validation and drop a stale coactivation comment in the composer. Signed-off-by: IRENEKO --- .../composer/compose_granite_switch.py | 1 - src/granite_switch/composer/compose_utils.py | 12 ++---------- src/granite_switch/vllm/core/lora_kernel_meta.py | 7 ------- 3 files changed, 2 insertions(+), 18 deletions(-) diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index 0882b044..65af332a 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -1302,7 +1302,6 @@ def build(): per_slot_label_token_ids.append( resolve_label_token_ids(tokenizer, label_names) ) - per_slot_read_layer.append(classifier_meta[name].get("layer")) optional_kwargs["adapter_kinds"] = adapter_kinds diff --git a/src/granite_switch/composer/compose_utils.py b/src/granite_switch/composer/compose_utils.py index a3bc87b9..ca3c2802 100644 --- a/src/granite_switch/composer/compose_utils.py +++ b/src/granite_switch/composer/compose_utils.py @@ -174,9 +174,7 @@ def from_base_and_adapters( # adapter_ranks/adapter_alphas stay full external width (one entry per # slot, classifiers included) so config.adapter_ranks is length - # num_adapters. A classifier slot itself has no LoRA weights. A - # separately selected LoRA slot can remain active through a classifier - # marker because classifier selection does not overwrite adapter state. + # num_adapters. A classifier slot itself has no LoRA weights. if lora_adapter_paths: lora_rank, _lora_alpha, lora_ranks, lora_alphas = detect_lora_config( lora_adapter_paths @@ -307,9 +305,7 @@ def from_base_and_adapters( base_model_name_or_path, model, switch_config, arch ) - adapter_mapping = {} classifier_mapping = {} - if lora_adapter_paths: # --- Step 7: Transfer LoRA adapter weights --- # LoRAs are the contiguous leading slots, so the leading k entries of @@ -330,11 +326,7 @@ def from_base_and_adapters( model, arch, adapter_paths=lora_adapter_paths, - adapter_names=( - [adapter_names[i] for i in lora_slots] - if adapter_names is not None - else None - ), + adapter_names=lora_adapter_names, target_module_sets=target_module_sets, ) if dual_stream: diff --git a/src/granite_switch/vllm/core/lora_kernel_meta.py b/src/granite_switch/vllm/core/lora_kernel_meta.py index 94f6ca6f..f7b740b9 100644 --- a/src/granite_switch/vllm/core/lora_kernel_meta.py +++ b/src/granite_switch/vllm/core/lora_kernel_meta.py @@ -192,7 +192,6 @@ class LoRAContext: __slots__ = ( "adapter_indices", - "classifier_indices", "num_tokens", "per_module_bitmasks", "remapped_indices", @@ -203,18 +202,12 @@ def __init__(self): self.remapped_indices: torch.Tensor | None = None self.per_module_bitmasks: torch.Tensor | None = None self.num_tokens: int = 0 - # Classifier stream (parallel to the LoRA stream): 0 = no classifier, - # 1+ = classifier slot. Set by the model's forward when a classifier - # head is present; read by the verdict exit on the last rank. Not part - # of the fused-LoRA bitmask path — the LoRA kernel never sees it. - self.classifier_indices: torch.Tensor | None = None def reset(self): self.adapter_indices = None self.remapped_indices = None self.per_module_bitmasks = None self.num_tokens = 0 - self.classifier_indices = None class FusedLoRAKernelMeta(nn.Module): From 548104dc66c03310baa4a8f425e9796123c39333 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Wed, 21 Oct 2026 22:44:45 +0000 Subject: [PATCH 18/20] refactor(classifier): shorten marker bookkeeping, head transfer and guide Move HF marker location into one helper that treats the final layer as another read layer. Drop the classifier mapping record nobody read, inline its one-use validator, reuse the shared head helpers in the vLLM worker, and tighten the guide's overview and internals. Signed-off-by: IRENEKO --- src/granite_switch/composer/compose_utils.py | 4 +- .../composer/weight_transfer.py | 63 ++----- .../hf/modeling_granite_switch.py | 165 +++++++++--------- .../_classifier_chunked_prefill_worker.py | 45 +---- tutorials/guides/classifier_slots.md | 43 ++--- 5 files changed, 117 insertions(+), 203 deletions(-) diff --git a/src/granite_switch/composer/compose_utils.py b/src/granite_switch/composer/compose_utils.py index ca3c2802..6a5f1191 100644 --- a/src/granite_switch/composer/compose_utils.py +++ b/src/granite_switch/composer/compose_utils.py @@ -305,7 +305,6 @@ def from_base_and_adapters( base_model_name_or_path, model, switch_config, arch ) - classifier_mapping = {} if lora_adapter_paths: # --- Step 7: Transfer LoRA adapter weights --- # LoRAs are the contiguous leading slots, so the leading k entries of @@ -338,7 +337,7 @@ def from_base_and_adapters( # --- Step 7b: Transfer trained classifier-head weights --- # Each slot's head is validated against its own label count and # copied into the top rows of the padded bank. - classifier_mapping = transfer_classifier_weights( + transfer_classifier_weights( classifier_adapter_paths, classifier_slots, model, @@ -359,7 +358,6 @@ def from_base_and_adapters( model._build_mappings = { "base": base_mapping, "adapter": adapter_mapping, - "classifier": classifier_mapping, "source_analysis": source_analysis, # Per-external-adapter alpha, parallel to adapter_paths. When no # external adapters are provided, detect_lora_config isn't called diff --git a/src/granite_switch/composer/weight_transfer.py b/src/granite_switch/composer/weight_transfer.py index f5f2472c..0713f498 100644 --- a/src/granite_switch/composer/weight_transfer.py +++ b/src/granite_switch/composer/weight_transfer.py @@ -842,8 +842,7 @@ def transfer_classifier_weights( model, num_labels_per_slot: list[int], hidden_size: int, - return_mapping: bool = True, -) -> dict | None: +) -> None: """Load trained classifier heads and transfer them into the switch model. A classifier slot is the classifier-head analog of a LoRA adapter: it shares @@ -863,10 +862,6 @@ def transfer_classifier_weights( (``config.classifier_num_labels_per_slot``, indexed by ``adapter_index - 1``). hidden_size: Expected hidden size (``config.hidden_size``). - return_mapping: If True, return detailed mapping information. - - Returns: - Mapping record dict, or None if *return_mapping* is False. """ from safetensors.torch import load_file @@ -880,9 +875,20 @@ def transfer_classifier_weights( "the head is constructed before loading trained weights." ) - _validate_classifier_transfer(model, slot_indices) - - mapping_record = {"source_params": [], "target_params": [], "mappings": []} + # Two directories on one bank row, or a row outside the bank, would + # silently corrupt another slot's head. + if len(slot_indices) != len(set(slot_indices)): + raise ValueError( + f"Duplicate classifier slot indices {slot_indices}; each classifier " + "must map to a distinct bank row." + ) + num_rows = head.weight.shape[0] + for slot in slot_indices: + if not (0 <= slot < num_rows): + raise ValueError( + f"Classifier slot {slot} is out of range for a classifier bank " + f"of {num_rows} rows (num_adapters)." + ) print(f"Loading {len(classifier_paths)} trained classifier head(s)...") with torch.no_grad(): @@ -931,44 +937,5 @@ def transfer_classifier_weights( ) print(f" Classifier: {path} -> classifier_head[{slot}, :{n_labels}]") - mapping_record["source_params"].append(f"{path}::weight") - mapping_record["source_params"].append(f"{path}::bias") - mapping_record["mappings"].append( - { - "source": [f"{path}::weight", f"{path}::bias"], - "target": f"model.classifier_head[{slot}]", - "type": "classifier", - } - ) - - mapping_record["target_params"] = [ - f"model.classifier_head.weight[{s}]" for s in slot_indices - ] + [f"model.classifier_head.bias[{s}]" for s in slot_indices] print(f"Loaded {len(classifier_paths)} classifier head(s)") - return mapping_record if return_mapping else None - - -def _validate_classifier_transfer(model, slot_indices: list[int]): - """Sanity-check the classifier slots before transfer. - - Guards against two silent-corruption modes: two directories claiming the - same bank row, and a slot index outside the head's ``[num_adapters, ...]`` - bank. - """ - if len(slot_indices) != len(set(slot_indices)): - raise ValueError( - f"Duplicate classifier slot indices {slot_indices}; each classifier " - "must map to a distinct bank row." - ) - num_rows = model.model.classifier_head.weight.shape[0] - for slot in slot_indices: - if not (0 <= slot < num_rows): - raise ValueError( - f"Classifier slot {slot} is out of range for a classifier bank " - f"of {num_rows} rows (num_adapters)." - ) - print( - f"Classifier transfer: {len(slot_indices)} slot(s) " - f"{sorted(slot_indices)} into a {num_rows}-row bank" - ) diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index a4d812ad..8f782ced 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -479,6 +479,59 @@ def __init__(self, config: GraniteSwitchConfig): for layer in self.layers: nn.init.zeros_(layer.cross_stream.base_layer.weight) + def _locate_classifier_markers(self, input_ids, classifier_indices, attention_mask): + """``(rows, pos, slots, markers_by_layer)`` for this forward's markers. + + ``markers_by_layer`` maps each read layer (None = final) to the marker + indices that read it. ``None`` when there is no classifier marker. Each + marker must be its row's last real token, since the verdict replaces + that position's logits. + """ + if self.classifier_head is None: + return None + if input_ids is None: + if bool((classifier_indices > 0).any()): + raise RuntimeError( + "Classifier requests need input_ids, not inputs_embeds." + ) + return None + marker_ids = torch.tensor( + self.config.classifier_control_token_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + rows, pos = torch.isin(input_ids, marker_ids).nonzero(as_tuple=True) + if rows.numel() == 0: + return None + + batch_size, seq_length = input_ids.shape + if attention_mask is not None and attention_mask.dim() == 2: + real = attention_mask[:, -seq_length:].long() + cols = torch.arange(seq_length, device=real.device) + last = (real * cols).max(dim=1).values + else: + last = torch.full((batch_size,), seq_length - 1, device=input_ids.device) + misplaced = pos != last[rows] + if bool(misplaced.any()): + bad = sorted(set(rows[misplaced].tolist())) + raise RuntimeError( + f"Classifier request(s) {bad}: control token is not their last token." + ) + + slots = classifier_indices[rows, pos] + read_layers = self.config.classifier_read_layers or [None] * len( + self.config.adapter_kinds + ) + by_layer: dict[int | None, list[int]] = {} + for marker_idx, slot in enumerate(slots.tolist()): + layer = read_layers[slot - 1] if slot > 0 else None + by_layer.setdefault(layer, []).append(marker_idx) + markers_by_layer = { + layer: torch.tensor(indices, device=input_ids.device) + for layer, indices in by_layer.items() + } + return rows, pos, slots, markers_by_layer + def forward( self, input_ids: torch.LongTensor | None = None, @@ -649,77 +702,29 @@ def forward( all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None - # Resolve classifier read positions before the decoder runs. This lets - # intermediate heads consume only their marker rows as each layer - # produces them, without retaining full sequence activations per layer. + # Resolve classifier read positions before the decoder runs, so each read + # layer scores only its marker rows as it produces them. + markers = self._locate_classifier_markers( + input_ids, classifier_indices, attention_mask + ) self._last_classifier_verdict = None - classifier_rows = classifier_pos = classifier_slots = None - classifier_verdict = None - classifier_markers_by_layer = {} - classifier_final_markers = None - if self.classifier_head is not None: - if input_ids is None: - if bool((classifier_indices > 0).any()): - raise RuntimeError( - "Classifier requests need input_ids, not inputs_embeds." - ) - else: - marker_ids = torch.tensor( - self.config.classifier_control_token_ids, - device=input_ids.device, - dtype=input_ids.dtype, - ) - is_marker = torch.isin(input_ids, marker_ids) - if bool(is_marker.any()): - classifier_rows, classifier_pos = is_marker.nonzero(as_tuple=True) - classifier_slots = classifier_indices[ - classifier_rows, classifier_pos - ] - - # Each marker must be the last real token of its request, - # since the verdict replaces that position's logits. - if attention_mask is not None and attention_mask.dim() == 2: - real = attention_mask[:, -seq_length:].long() - cols = torch.arange(seq_length, device=real.device) - last = (real * cols).max(dim=1).values - else: - last = torch.full( - (batch_size,), seq_length - 1, device=input_ids.device - ) - misplaced = classifier_pos != last[classifier_rows] - if bool(misplaced.any()): - bad = sorted(set(classifier_rows[misplaced].tolist())) - raise RuntimeError( - f"Classifier request(s) {bad}: control token is not " - "their last token." - ) - - read_layers = self.config.classifier_read_layers - if read_layers is None: - read_layers = [None] * self.config.num_adapters - layer_to_marker_indices = {} - final_marker_indices = [] - for marker_idx, slot in enumerate(classifier_slots.tolist()): - layer = read_layers[slot - 1] if slot > 0 else None - if layer is None: - final_marker_indices.append(marker_idx) - else: - layer_to_marker_indices.setdefault(layer, []).append( - marker_idx - ) - classifier_markers_by_layer = { - layer: torch.tensor(indices, device=input_ids.device) - for layer, indices in layer_to_marker_indices.items() - } - classifier_final_markers = torch.tensor( - final_marker_indices, device=input_ids.device - ) - classifier_verdict = inputs_embeds.new_zeros( - ( - classifier_rows.numel(), - self.config.max_classifier_labels, - ) - ) + if markers is not None: + classifier_rows, classifier_pos, classifier_slots, markers_by_layer = ( + markers + ) + classifier_verdict = inputs_embeds.new_zeros( + (classifier_rows.numel(), self.config.max_classifier_labels) + ) + + def score_markers(layer, hidden): + """Score the markers that read ``layer`` (None = final) from ``hidden``.""" + if markers is None or layer not in markers_by_layer: + return + idx = markers_by_layer[layer] + classifier_verdict[idx] = self.classifier_head( + hidden[classifier_rows[idx], classifier_pos[idx]], + classifier_slots[idx], + ) for layer_idx, decoder_layer in enumerate(self.layers): if output_hidden_states: @@ -743,14 +748,7 @@ def forward( else: hidden_states = layer_outputs[0] - marker_indices = classifier_markers_by_layer.get(layer_idx) - if marker_indices is not None: - marker_hidden = hidden_states[ - classifier_rows[marker_indices], classifier_pos[marker_indices] - ] - classifier_verdict[marker_indices] = self.classifier_head( - marker_hidden, classifier_slots[marker_indices] - ) + score_markers(layer_idx, hidden_states) if output_attentions: if layer_outputs[1] is not None: @@ -761,15 +759,8 @@ def forward( if output_hidden_states: all_hidden_states += (hidden_states,) - if classifier_rows is not None: - if classifier_final_markers.numel() > 0: - final_hidden = hidden_states[ - classifier_rows[classifier_final_markers], - classifier_pos[classifier_final_markers], - ] - classifier_verdict[classifier_final_markers] = self.classifier_head( - final_hidden, classifier_slots[classifier_final_markers] - ) + if markers is not None: + score_markers(None, hidden_states) self._last_classifier_verdict = ( classifier_rows, classifier_pos, diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_classifier_chunked_prefill_worker.py index d60f8953..c888c9d3 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_classifier_chunked_prefill_worker.py @@ -41,47 +41,18 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): ``num_labels`` single-token labels are taken from ``_LABEL_CANDIDATES`` and ``scale`` is the std of the random head weight. """ - import torch - from safetensors.torch import save_file from transformers import AutoConfig, AutoTokenizer from granite_switch.composer.compose_utils import GraniteSwitchComposer - from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE tok = AutoTokenizer.from_pretrained(base) - hidden = AutoConfig.from_pretrained(base).hidden_size - - labels, ids = [], [] - for c in _LABEL_CANDIDATES: - e = tok.encode(c, add_special_tokens=False) - if len(e) == 1 and e[0] not in ids: - labels.append(c) - ids.append(e[0]) - if len(labels) == num_labels: - break - if len(labels) < num_labels: - raise RuntimeError( - f"only {len(labels)} single-token labels available from " - f"{_LABEL_CANDIDATES}, need {num_labels}" - ) - nl = len(labels) - - g = torch.Generator().manual_seed(0) + labels, ids = _pick_labels(tok, num_labels) # Weight-based (not bias-only) head so the verdict depends on the hidden state. - weight = torch.randn((nl, hidden), generator=g, dtype=torch.float32) * scale - bias = torch.zeros((nl,), dtype=torch.float32) - - slot = os.path.join(outdir, "detect") - os.makedirs(slot, exist_ok=True) - save_file( - {"weight": weight.contiguous(), "bias": bias.contiguous()}, - os.path.join(slot, CLASSIFIER_HEAD_FILE), - ) - json.dump( - {"kind": "classifier", "labels": labels}, - open(os.path.join(slot, "adapter_config.json"), "w"), + weight, bias = _fixed_head( + 0, len(ids), AutoConfig.from_pretrained(base).hidden_size, scale ) - open(os.path.join(slot, "io.yaml"), "w").write("name: detect\nmodel: ~\n") + slot = os.path.join(outdir, "detect") + _write_head(slot, weight, bias, labels) control_id = tok.vocab_size - control_offset model = GraniteSwitchComposer.from_base_and_adapters( @@ -168,13 +139,13 @@ def _pick_labels(tok, num_labels): raise RuntimeError(f"only {len(labels)} single-token labels, need {num_labels}") -def _fixed_head(seed, num_labels, hidden_size): - """Fixed seed-``seed`` weight (std 0.5) and zero bias. Do not search seeds.""" +def _fixed_head(seed, num_labels, hidden_size, scale=0.5): + """Fixed seed-``seed`` weight (std ``scale``) and zero bias. Do not search seeds.""" import torch g = torch.Generator().manual_seed(seed) weight = torch.randn((num_labels, hidden_size), generator=g, dtype=torch.float32) - return weight * 0.5, torch.zeros(num_labels, dtype=torch.float32) + return weight * scale, torch.zeros(num_labels, dtype=torch.float32) def _write_head(slot_dir, weight, bias, labels): diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index fed620d1..2cfbcb46 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -6,15 +6,11 @@ serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detecti ## Overview -A Granite Switch adapter slot is normally a LoRA adapter. A **classifier slot** is an alternative -occupant of that same slot: instead of a LoRA, it holds a small linear classification head. It is -triggered by the *same* control-token machinery — one `<|name|>` token, one adapter index — but -the classifier reads the configured layer's activation and produces per-label logits. - -Activation probes are supported as classifier slots in the standard composition workflow. Probes and -generation adapters can coexist in one checkpoint and be selected individually: a request selects -one slot, either a generation adapter or a classifier. The calling stack manages compatible -execution state and cache reuse. +A **classifier slot** holds a small linear classification head instead of a LoRA. It uses the same +control-token machinery (one `<|name|>` token, one adapter index) and reads the configured layer's +activation to produce per-label logits. Classifier and generation-adapter slots can coexist in one +checkpoint; a request selects one slot. The calling stack manages compatible execution state and +cache reuse. At inference the verdict does not come back as a raw score vector. Each label word is resolved to a single token id at compose time, and the model rewrites the request's final logit row so that the @@ -334,25 +330,16 @@ channel. ## How it works (brief) -- **Same trigger, separate path.** One `<|name|>` control token selects a classifier slot, exactly - like an adapter. Classifier positions are routed out of the LoRA stream, and the head reads the - activation at that token. -- **Padded head bank.** All classifier heads live in one stacked tensor padded to - `max_classifier_labels`; a slot with fewer labels leaves its extra rows zero. This keeps the vLLM - head a single fused, compile-safe operation regardless of how many labels each slot has — the same - way the LoRA bank pads every adapter to `max_lora_rank`. -- **Scored at the read layer.** A slot reading an intermediate layer is scored as soon as that - layer runs; final-layer slots are scored after the final norm. The per-label scores and slot id - travel with the hidden states, so vLLM's own row selection delivers them to the logits step, - which also works under CUDA-graph capture. -- **Pipeline transport.** The first stage identifies marker tokens. Each stage scores its own - configured layers and passes those scores and marker indices to the next stage; intermediate - hidden states do not need to be saved for classification on the last stage. -- **Label-word verdict exit.** At each classifier request's final token the model blanks the - vocabulary and scatters the head's per-label logits onto that slot's own label token ids (sliced - to the slot's real label count, so padded columns are never emitted). The sampler then emits the - label word. - +- **Routing.** A classifier control token selects its slot like any adapter; its position is routed + out of the LoRA stream and the head reads the activation there. +- **Padded head bank.** All heads share one tensor padded to `max_classifier_labels` (like the LoRA + bank pads to `max_lora_rank`), keeping the vLLM head one compile-safe operation. +- **Scoring.** Intermediate-layer slots are scored as that layer runs, final-layer slots after the + final norm. Scores and slot id travel with the hidden states to the logits step, so this also + works under CUDA-graph capture. Under pipeline parallelism the first stage finds the markers and + each stage passes scores and marker indices on. +- **Verdict exit.** The marker's logit row is blanked and the head's scores are written onto that + slot's label token ids, so the sampler emits the label word. ## Related From 731b368e9d648b3449aad273c98aa9b8e78ec169 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Thu, 22 Oct 2026 19:42:06 +0000 Subject: [PATCH 19/20] docs(classifier): drop the cache-reuse note from the guide overview The guide covers composing and selecting classifier slots; execution state and cache reuse are outside that scope. Signed-off-by: IRENEKO --- tutorials/guides/classifier_slots.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/classifier_slots.md index 2cfbcb46..291c68df 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/classifier_slots.md @@ -9,8 +9,7 @@ serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detecti A **classifier slot** holds a small linear classification head instead of a LoRA. It uses the same control-token machinery (one `<|name|>` token, one adapter index) and reads the configured layer's activation to produce per-label logits. Classifier and generation-adapter slots can coexist in one -checkpoint; a request selects one slot. The calling stack manages compatible execution state and -cache reuse. +checkpoint; a request selects one slot. At inference the verdict does not come back as a raw score vector. Each label word is resolved to a single token id at compose time, and the model rewrites the request's final logit row so that the From 7ca9aff5bf25711012ec578903770e12b3dc2e98 Mon Sep 17 00:00:00 2001 From: IRENEKO Date: Fri, 23 Oct 2026 01:32:28 +0000 Subject: [PATCH 20/20] refactor(probe): rename classifier slots to activation probes Identifiers, module and test paths, config fields, the manifest kind, the head file and the guide now say probe (activation probe in prose). Checkpoints composed under the classifier names need re-composing; the feature has not shipped on main. Signed-off-by: IRENEKO --- .../composer/adapter_discovery.py | 18 +- src/granite_switch/composer/arch.py | 4 +- .../composer/compose_granite_switch.py | 60 +++--- src/granite_switch/composer/compose_utils.py | 54 +++--- .../composer/reporting/population_table.py | 2 +- .../composer/tokenizer_setup.py | 46 ++--- src/granite_switch/composer/validator.py | 2 +- .../composer/weight_transfer.py | 66 +++---- src/granite_switch/config.py | 89 ++++----- src/granite_switch/hf/core/__init__.py | 4 +- .../hf/core/{classifier.py => probe.py} | 43 +++-- .../hf/modeling_granite_switch.py | 100 +++++----- src/granite_switch/hf/switch/multi.py | 10 +- ...classifier_routing.py => probe_routing.py} | 36 ++-- src/granite_switch/vllm/core/__init__.py | 4 +- .../vllm/core/{classifier.py => probe.py} | 43 +++-- src/granite_switch/vllm/decoder/interface.py | 2 +- .../vllm/granite_switch_model.py | 178 +++++++++-------- src/granite_switch/vllm/switch/multi.py | 16 +- tests/composer/test_adapter_loader.py | 26 +-- tests/composer/test_base_reset_token.py | 10 +- tests/composer/test_control_lut_refresh.py | 16 +- .../composer/test_multi_audio_compose_e2e.py | 46 +++-- ...ifier_compose.py => test_probe_compose.py} | 134 +++++++------ tests/composer/test_sr_compose_e2e.py | 34 ++-- tests/composer/test_tokenizer_setup.py | 14 +- .../hf/{test_classifier.py => test_probe.py} | 180 +++++++++--------- ...assifier_recorder.py => probe_recorder.py} | 16 +- tests/unit/test_config.py | 50 +++-- ...ifier_routing.py => test_probe_routing.py} | 38 ++-- tests/vllm/_model_forward_tests.py | 18 +- ...er.py => _probe_chunked_prefill_worker.py} | 80 ++++---- ...safety.py => test_probe_capture_safety.py} | 84 ++++---- ...efill.py => test_probe_chunked_prefill.py} | 36 ++-- tutorials/README.md | 2 +- ...ier_slots.md => activation_probe_slots.md} | 107 +++++------ 36 files changed, 815 insertions(+), 853 deletions(-) rename src/granite_switch/hf/core/{classifier.py => probe.py} (69%) rename src/granite_switch/{classifier_routing.py => probe_routing.py} (50%) rename src/granite_switch/vllm/core/{classifier.py => probe.py} (62%) rename tests/composer/{test_classifier_compose.py => test_probe_compose.py} (79%) rename tests/hf/{test_classifier.py => test_probe.py} (65%) rename tests/shared/{classifier_recorder.py => probe_recorder.py} (81%) rename tests/unit/{test_classifier_routing.py => test_probe_routing.py} (61%) rename tests/vllm/{_classifier_chunked_prefill_worker.py => _probe_chunked_prefill_worker.py} (91%) rename tests/vllm/{test_classifier_capture_safety.py => test_probe_capture_safety.py} (80%) rename tests/vllm/{test_classifier_chunked_prefill.py => test_probe_chunked_prefill.py} (94%) rename tutorials/guides/{classifier_slots.md => activation_probe_slots.md} (75%) diff --git a/src/granite_switch/composer/adapter_discovery.py b/src/granite_switch/composer/adapter_discovery.py index eb5b334e..9fd2cd83 100644 --- a/src/granite_switch/composer/adapter_discovery.py +++ b/src/granite_switch/composer/adapter_discovery.py @@ -112,20 +112,20 @@ def discover_adapters_from_yaml( """Discover adapters from a YAML manifest file. Reads a YAML manifest that maps adapter names to their paths and types. - An entry can also mark itself a **classifier** slot with these extra keys: + An entry can also mark itself a **probe** slot with these extra keys: .. code-block:: yaml hallucination: path: /path/to/head type: lora # lora/alora technology (token placement) - kind: classifier # default "lora" when omitted + kind: probe # default "lora" when omitted labels: [faithful, hallucinated] layer: 12 # optional decoder layer; omit for the final layer ``type`` and ``kind`` are independent axes: ``type`` is the LoRA technology (lora vs alora, which drives token placement), ``kind`` is whether the slot - is a real LoRA or a classifier head. ``kind`` and ``labels`` are optional; + is a real LoRA or a probe head. ``kind`` and ``labels`` are optional; without them an entry is a plain LoRA adapter. The label words become token ids at compose time (see :func:`tokenizer_setup.resolve_label_token_ids`). @@ -136,10 +136,10 @@ def discover_adapters_from_yaml( manifest_path: Path to the YAML manifest file. Returns: - ``(found, classifier_meta)``. ``found`` is: + ``(found, probe_meta)``. ``found`` is: List of ``(adapter_path, adapter_name, technology, source)`` tuples. The source is set to the manifest path for traceability. - ``classifier_meta`` maps name -> + ``probe_meta`` maps name -> ``{"kind", "labels", "layer"}`` for each entry that set any of them; names not in it are plain LoRA slots. """ @@ -149,7 +149,7 @@ def discover_adapters_from_yaml( print(f" Loading adapters manifest: {path.name}") found = [] - classifier_meta: dict[str, dict] = {} + probe_meta: dict[str, dict] = {} if path.is_file() and path.suffix in (".yaml", ".yml"): with open(path) as f: adapters_config = yaml.safe_load(f) @@ -161,18 +161,18 @@ def discover_adapters_from_yaml( # Use manifest path as source for traceability found.append((adapter_path, name, tech, manifest_path)) - # Optional classifier slots. + # Optional probe slots. kind = info.get("kind") labels = info.get("labels") layer = info.get("layer") if kind is not None or labels is not None or layer is not None: - classifier_meta[name] = { + probe_meta[name] = { "kind": kind or "lora", "labels": labels, "layer": layer, } - return found, classifier_meta + return found, probe_meta def _report_module_contributions( diff --git a/src/granite_switch/composer/arch.py b/src/granite_switch/composer/arch.py index c5b46100..fa0ff8b5 100644 --- a/src/granite_switch/composer/arch.py +++ b/src/granite_switch/composer/arch.py @@ -112,13 +112,13 @@ class ArchDescriptor: # LoRA parameter keywords (for filtering) lora_keywords: list[str] = field(default_factory=lambda: ["lora_A", "lora_B"]) - # Non-LoRA/classifier buffer keywords to exclude from base validation + # Non-LoRA/probe buffer keywords to exclude from base validation buffer_keywords: list[str] = field( default_factory=lambda: [ "adapter_token_ids", "adapter_scalings", "control_to_substitute_lut", - "classifier_head", + "probe_head", "adapter_kind_lut", ] ) diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index 65af332a..d7971982 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -69,7 +69,7 @@ from granite_switch.composer.compose_utils import GraniteSwitchComposer from granite_switch.composer.reporting import generate_compose_report, write_build_doc from granite_switch.composer.tokenizer_setup import ( - ANCHOR_MODE_CLASSIFIER, + ANCHOR_MODE_PROBE, ANCHOR_MODE_SR, add_audio_token, add_control_tokens, @@ -122,17 +122,17 @@ def build_control_token_lists(tokenizer, all_discovered, base_reset): all_discovered, _probe_lora_substitute_token_id(tokenizer), base_reset=base_reset, - classifier_substitute_id=_classifier_substitute_token_id(tokenizer), + probe_slot_substitute_id=_probe_slot_substitute_token_id(tokenizer), ) return adapter_token_ids, special_tokens, adapter_substitute_token_ids -def _classifier_substitute_token_id(tokenizer) -> int: +def _probe_slot_substitute_token_id(tokenizer) -> int: """The base's role-open marker id, which follows the last turn's close.""" fmt = detect_template_format(tokenizer.chat_template) if fmt is None: raise ValueError( - "Unrecognized chat template; cannot pick the classifier substitute." + "Unrecognized chat template; cannot pick the probe substitute." ) ids = tokenizer.encode(fmt.role_open_marker, add_special_tokens=False) if len(ids) != 1: @@ -1000,7 +1000,7 @@ def build(): step_start = time.time() discovered_adapters = [] - classifier_meta: dict[str, dict] = {} + probe_meta: dict[str, dict] = {} if args.adapters: print("\n" + "=" * 80) print("Processing adapters") @@ -1042,9 +1042,9 @@ def build(): print(msg) discovered_adapters.extend(found) elif (path := Path(entry)).is_file() and path.suffix in (".yaml", ".yml"): - found, manifest_classifier_meta = discover_adapters_from_yaml(entry) + found, manifest_probe_meta = discover_adapters_from_yaml(entry) discovered_adapters.extend(found) - classifier_meta.update(manifest_classifier_meta) + probe_meta.update(manifest_probe_meta) else: # Single adapter directory resolved = Path(resolved_path) @@ -1086,11 +1086,11 @@ def build(): # token lands, and SR has exactly one usable activation point — the anchor at # the end of the generation prompt — so a mislabelled SR checkpoint must not # be able to fall through to aLoRA or sequence-start placement. - # A classifier slot is labelled here for the same reason: its control token is + # A probe slot is labelled here for the same reason: its control token is # an end-locator, which is a placement mode, not a weight format, and the # manifest's ``kind`` decides it. An adapter satisfies at most one test -- an SR - # adapter carries ``cross_stream`` weights, a classifier slot carries a - # classifier head instead -- and SR goes first because it is read from the + # adapter carries ``cross_stream`` weights, a probe slot carries a + # probe head instead -- and SR goes first because it is read from the # artifact rather than from a label. external_discovered = [ ( @@ -1098,8 +1098,8 @@ def build(): name, ANCHOR_MODE_SR if path and is_shadow_residual_adapter(path) - else ANCHOR_MODE_CLASSIFIER - if classifier_meta.get(name, {}).get("kind") == "classifier" + else ANCHOR_MODE_PROBE + if probe_meta.get(name, {}).get("kind") == "probe" else tech, source, ) @@ -1110,9 +1110,9 @@ def build(): ] all_discovered = external_discovered + built_in_discovered - # Order LoRA slots first, classifier slots last (Python sorts False before True) + # Order LoRA slots first, probe slots last (Python sorts False before True) all_discovered.sort( - key=lambda t: classifier_meta.get(t[1], {}).get("kind", "lora") == "classifier" + key=lambda t: probe_meta.get(t[1], {}).get("kind", "lora") == "probe" ) has_external = len(external_discovered) > 0 @@ -1263,37 +1263,37 @@ def build(): _anchor_text, _sub, _mode = load_activation_anchor(_path, tokenizer) adapter_substitute_token_ids[_sub_offset + _i] = _sub - # adapter_kinds marks which slots are classifiers (same control token as a + # adapter_kinds marks which slots are probes (same control token as a # LoRA); resolve each slot's label words to token ids for the vLLM verdict. - if classifier_meta: + if probe_meta: adapter_kinds = [ - classifier_meta.get(name, {}).get("kind", "lora") for name in adapter_names + probe_meta.get(name, {}).get("kind", "lora") for name in adapter_names ] - classifier_names = [ + probe_names = [ name for name in adapter_names - if classifier_meta.get(name, {}).get("kind") == "classifier" + if probe_meta.get(name, {}).get("kind") == "probe" ] - if classifier_names: + if probe_names: # Full width (length num_adapters, indexed by adapter_index - 1): each - # classifier slot's label token ids. The config derives the per-slot + # probe slot's label token ids. The config derives the per-slot # counts and the padded bank width from this. per_slot_label_names: list[list[str] | None] = [] per_slot_label_token_ids: list[list[int] | None] = [] - # Full width too: each classifier slot's read layer (None = final). + # Full width too: each probe slot's read layer (None = final). per_slot_read_layer: list[int | None] = [] for name, kind in zip(adapter_names, adapter_kinds): - if kind != "classifier": + if kind != "probe": per_slot_label_names.append(None) per_slot_label_token_ids.append(None) per_slot_read_layer.append(None) continue - labels = classifier_meta[name].get("labels") + labels = probe_meta[name].get("labels") if not labels: raise ValueError( - f"Adapter '{name}' is declared kind=classifier but has no " - f"'labels' in the manifest. Classifier slots need label " + f"Adapter '{name}' is declared kind=probe but has no " + f"'labels' in the manifest. Probe slots need label " f"words (e.g. labels: [safe, unsafe]) so the verdict exit " f"can emit them." ) @@ -1302,12 +1302,12 @@ def build(): per_slot_label_token_ids.append( resolve_label_token_ids(tokenizer, label_names) ) - per_slot_read_layer.append(classifier_meta[name].get("layer")) + per_slot_read_layer.append(probe_meta[name].get("layer")) optional_kwargs["adapter_kinds"] = adapter_kinds - optional_kwargs["classifier_label_token_ids"] = per_slot_label_token_ids - optional_kwargs["classifier_read_layers"] = per_slot_read_layer - print("\nClassifier slots:") + optional_kwargs["probe_label_token_ids"] = per_slot_label_token_ids + optional_kwargs["probe_read_layers"] = per_slot_read_layer + print("\nProbe slots:") for name, words, ids, layer in zip( adapter_names, per_slot_label_names, diff --git a/src/granite_switch/composer/compose_utils.py b/src/granite_switch/composer/compose_utils.py index 6a5f1191..75642f8e 100644 --- a/src/granite_switch/composer/compose_utils.py +++ b/src/granite_switch/composer/compose_utils.py @@ -20,7 +20,7 @@ from .weight_transfer import ( transfer_adapter_weights, transfer_base_weights, - transfer_classifier_weights, + transfer_probe_weights, ) @@ -90,23 +90,21 @@ def from_base_and_adapters( ext_kinds = ["lora"] * num_external else: ext_kinds = list(ext_kinds)[:num_external] - lora_slots = [i for i in range(num_external) if ext_kinds[i] != "classifier"] - classifier_slots = [ - i for i in range(num_external) if ext_kinds[i] == "classifier" - ] - - # Compose orders LoRAs first, classifiers last, so classifier slots must - # form a contiguous suffix (LoRAs occupy 0..k-1, classifiers k..). - if classifier_slots and classifier_slots != list( + lora_slots = [i for i in range(num_external) if ext_kinds[i] != "probe"] + probe_slots = [i for i in range(num_external) if ext_kinds[i] == "probe"] + + # Compose orders LoRAs first, probes last, so probe slots must + # form a contiguous suffix (LoRAs occupy 0..k-1, probes k..). + if probe_slots and probe_slots != list( range(lora_slots[-1] + 1 if lora_slots else 0, num_external) ): raise ValueError( - "Classifier slots must form a contiguous suffix (all LoRA slots " - f"first, then classifiers); got kinds {ext_kinds}. Order the " + "Probe slots must form a contiguous suffix (all LoRA slots " + f"first, then probes); got kinds {ext_kinds}. Order the " "adapters LoRAs-first (compose does this automatically)." ) lora_adapter_paths = [adapter_paths[i] for i in lora_slots] - classifier_adapter_paths = [adapter_paths[i] for i in classifier_slots] + probe_adapter_paths = [adapter_paths[i] for i in probe_slots] lora_adapter_names = ( [adapter_names[i] for i in lora_slots] if adapter_names is not None @@ -119,7 +117,7 @@ def from_base_and_adapters( base_config = load_base_config(base_model_name_or_path) # Classify adapting slots as single-stream (LoRA/aLoRA) or dual-stream - # (Shadow Residual). Classifier heads can read either decoder mode. + # (Shadow Residual). Probe heads can read either decoder mode. sr_adapters = [] non_sr_adapters = [] for ap in lora_adapter_paths: @@ -173,8 +171,8 @@ def from_base_and_adapters( ) # adapter_ranks/adapter_alphas stay full external width (one entry per - # slot, classifiers included) so config.adapter_ranks is length - # num_adapters. A classifier slot itself has no LoRA weights. + # slot, probes included) so config.adapter_ranks is length + # num_adapters. A probe slot itself has no LoRA weights. if lora_adapter_paths: lora_rank, _lora_alpha, lora_ranks, lora_alphas = detect_lora_config( lora_adapter_paths @@ -185,13 +183,13 @@ def from_base_and_adapters( adapter_names=lora_adapter_names, ) - # LoRAs lead, classifiers follow, so the full lists are a concat: - # detected LoRA values, then a placeholder per classifier slot. The + # LoRAs lead, probes follow, so the full lists are a concat: + # detected LoRA values, then a placeholder per probe slot. The # placeholder is the detected max rank, keeping max(adapter_ranks) == # max_lora_rank. - num_ext_classifier = num_external - len(lora_slots) - adapter_ranks = list(lora_ranks) + [lora_rank] * num_ext_classifier - adapter_alphas = list(lora_alphas) + [float(lora_rank)] * num_ext_classifier + num_ext_probe = num_external - len(lora_slots) + adapter_ranks = list(lora_ranks) + [lora_rank] * num_ext_probe + adapter_alphas = list(lora_alphas) + [float(lora_rank)] * num_ext_probe # Extend adapter_ranks with built-in entries if num_built_in > 0: @@ -205,8 +203,8 @@ def from_base_and_adapters( list(adapter_ranks) + [built_in_lora_rank] * num_built_in ) lora_rank = max(lora_rank, built_in_lora_rank) - elif classifier_adapter_paths or num_built_in > 0: - # No LoRA slots but the model still has slots (classifier and/or + elif probe_adapter_paths or num_built_in > 0: + # No LoRA slots but the model still has slots (probe and/or # built-in). Give every slot a placeholder LoRA rank so the # (unused) LoRA bank is well-formed and config validation passes. lora_rank = built_in_lora_rank @@ -333,15 +331,15 @@ def from_base_and_adapters( else: adapter_mapping = {} - if classifier_adapter_paths: - # --- Step 7b: Transfer trained classifier-head weights --- + if probe_adapter_paths: + # --- Step 7b: Transfer trained probe-head weights --- # Each slot's head is validated against its own label count and # copied into the top rows of the padded bank. - transfer_classifier_weights( - classifier_adapter_paths, - classifier_slots, + transfer_probe_weights( + probe_adapter_paths, + probe_slots, model, - num_labels_per_slot=switch_config.classifier_num_labels_per_slot, + num_labels_per_slot=switch_config.probe_num_labels_per_slot, hidden_size=switch_config.hidden_size, ) diff --git a/src/granite_switch/composer/reporting/population_table.py b/src/granite_switch/composer/reporting/population_table.py index 5cfec913..e19fc263 100644 --- a/src/granite_switch/composer/reporting/population_table.py +++ b/src/granite_switch/composer/reporting/population_table.py @@ -49,7 +49,7 @@ def generate_adapter_population_table( if target_module_sets is None: target_module_sets = load_adapter_target_modules(adapter_paths) - # target_module_sets covers LoRA slots only (classifier slots carry no LoRA + # target_module_sets covers LoRA slots only (probe slots carry no LoRA # weights, so source analysis never produced a target set for them). adapter_configs = [ { diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index 7dc83dee..a127835d 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -146,8 +146,8 @@ def detect_template_format(chat_template: str | None) -> TemplateFormat | None: ANCHOR_MODE_SR = "sr" #: Placement mode where the control token is emitted as the very last token, -#: after the last turn's close (classifier slot, manifest ``kind``). -ANCHOR_MODE_CLASSIFIER = "classifier" +#: after the last turn's close (probe slot, manifest ``kind``). +ANCHOR_MODE_PROBE = "probe" def _load_alora_invocation_token_ids(adapter_path: str) -> list[int]: @@ -836,7 +836,7 @@ def build_substitute_token_ids( discovered_adapters: list[tuple[str | None, str, str, str | None]], lora_substitute_id: int, base_reset: bool = False, - classifier_substitute_id: int | None = None, + probe_slot_substitute_id: int | None = None, ) -> list[int]: """Build the token-exchange substitute ids, index-aligned with the control tokens. @@ -850,14 +850,14 @@ def build_substitute_token_ids( * base-reset: ``lora_substitute_id`` as well. The token is placed at a turn boundary, where the role-open marker is exactly what the decoder expects at that position. - * classifier: ``classifier_substitute_id``, the base's role-open marker. + * probe: ``probe_slot_substitute_id``, the base's role-open marker. Args: discovered_adapters: ``(adapter_path, adapter_name, technology, source)`` tuples. lora_substitute_id: Probed sequence-start token id. base_reset: When True, prepend the base-reset slot's substitute so the list stays aligned with ``add_control_tokens(base_reset=True)``. - classifier_substitute_id: The base's role-open marker id. + probe_slot_substitute_id: The base's role-open marker id. Returns: One substitute id per control token, in the same order. @@ -866,17 +866,17 @@ def build_substitute_token_ids( for adapter_path, _name, technology, _source in discovered_adapters: if technology == "alora": substitute_ids.append(get_alora_first_invocation_token_id(adapter_path)) - elif technology == "classifier": - if classifier_substitute_id is None: - raise ValueError("A classifier slot needs classifier_substitute_id.") - substitute_ids.append(classifier_substitute_id) + elif technology == "probe": + if probe_slot_substitute_id is None: + raise ValueError("A probe slot needs probe_slot_substitute_id.") + substitute_ids.append(probe_slot_substitute_id) else: substitute_ids.append(lora_substitute_id) return substitute_ids def resolve_label_token_ids(tokenizer, label_names: list[str]) -> list[int]: - """Turn each classifier label word into its single token id.""" + """Turn each probe label word into its single token id.""" token_ids: list[int] = [] non_atomic: list[str] = [] for word in label_names: @@ -890,8 +890,8 @@ def resolve_label_token_ids(tokenizer, label_names: list[str]) -> list[int]: if non_atomic: raise ValueError( - f"{len(non_atomic)} classifier label word(s) do not encode to a " - f"single token, but the vLLM classifier exit writes the verdict at " + f"{len(non_atomic)} probe label word(s) do not encode to a " + f"single token, but the vLLM probe exit writes the verdict at " f"one token id per label. Choose a single-token synonym for each:\n " + "\n ".join(non_atomic) ) @@ -970,11 +970,11 @@ def configure_chat_template( # still costing compute. Callers therefore classify SR from the checkpoint # weights (``is_shadow_residual_adapter``) and pass ``sr`` here, rather than # trusting a directory label that no SR training run produces. - # A classifier slot arrives as technology ``classifier``, like SR arrives as + # A probe slot arrives as technology ``probe``, like SR arrives as # ``sr``: the caller derives the placement mode (SR from the checkpoint - # weights, classifier from the manifest's ``kind``) and passes it in + # weights, probe from the manifest's ``kind``) and passes it in # the discovered tuple. ``technology`` is the placement axis this function - # dispatches on; a classifier is a fourth mode alongside lora / alora / sr. + # dispatches on; a probe is a fourth mode alongside lora / alora / sr. adapter_mapping: dict[str, dict[str, str]] = {} sr_anchors: set[str] = set() for adapter_info in discovered_adapters: @@ -1025,7 +1025,7 @@ def configure_chat_template( for adapter_name, info in adapter_mapping.items(): # ``type`` is the placement mode: lora (sequence start) / alora (before # the invocation text) / sr (replaces the generation-prompt anchor) / - # classifier (very last token, after the last turn's close). + # probe (very last token, after the last turn's close). if "invocation_text" in info: mapping_entries.append( f" '{adapter_name}': {{'token': '{info['token']}', " @@ -1054,10 +1054,10 @@ def configure_chat_template( {%- set adapter_invocation_drop_chars = adapter_map[adapter_name]['invocation_drop_chars'] %} {%- endif %} {%- endif %} -{#- A classifier reads its verdict at the marker, so nothing may follow it: +{#- A probe reads its verdict at the marker, so nothing may follow it: drop the generation prompt even when the caller asks for one. -#} -{%- set classifier_mode = adapter_type == 'classifier' %} -{%- if classifier_mode %} +{%- set probe_mode = adapter_type == 'probe' %} +{%- if probe_mode %} {%- set add_generation_prompt = false %} {%- endif %} @@ -1331,11 +1331,11 @@ def configure_chat_template( + modified_chat_template[insertion_point:] ) - # Classifier placement: the marker is the very last token, after the last + # Probe placement: the marker is the very last token, after the last # turn's close. - if any(i["type"] == ANCHOR_MODE_CLASSIFIER for i in adapter_mapping.values()): + if any(i["type"] == ANCHOR_MODE_PROBE for i in adapter_mapping.values()): modified_chat_template += ( - "\n{%- if classifier_mode %}{{- adapter_token }}{%- endif %}" + "\n{%- if probe_mode %}{{- adapter_token }}{%- endif %}" ) # Insert alora fallback before generation prompt @@ -1441,7 +1441,7 @@ def _split_merged(match: "re.Match") -> str: tokenizer.chat_template = modified_chat_template print(f"Chat template configured with {len(adapter_mapping)} adapter mappings:") for adapter_name, info in adapter_mapping.items(): - if info["type"] == ANCHOR_MODE_CLASSIFIER: + if info["type"] == ANCHOR_MODE_PROBE: placement = "after the last turn's close (last token; read point = marker)" elif info["type"] == ANCHOR_MODE_SR: placement = f"replacing '{sr_site}' in the generation prompt" diff --git a/src/granite_switch/composer/validator.py b/src/granite_switch/composer/validator.py index 4bba80ff..a6c7b455 100644 --- a/src/granite_switch/composer/validator.py +++ b/src/granite_switch/composer/validator.py @@ -85,7 +85,7 @@ def validate_cross_stream_population(model): # slots in a freshly built model, but lora_B decides whether the branch # contributes anything. for adapter_idx in range(param.shape[0]): - if model.config.adapter_kinds[adapter_idx] == "classifier": + if model.config.adapter_kinds[adapter_idx] == "probe": continue if bool(torch.all(param[adapter_idx] == 0)): empty.append(f"{name}[{adapter_idx}]") diff --git a/src/granite_switch/composer/weight_transfer.py b/src/granite_switch/composer/weight_transfer.py index 0713f498..adebda4c 100644 --- a/src/granite_switch/composer/weight_transfer.py +++ b/src/granite_switch/composer/weight_transfer.py @@ -15,10 +15,10 @@ from .arch import ArchDescriptor -# Filename + keys of a trained classifier-head artifact inside a classifier +# Filename + keys of a trained probe-head artifact inside a probe # adapter directory. -CLASSIFIER_HEAD_FILE = "classifier_head.safetensors" -CLASSIFIER_HEAD_KEYS = ("weight", "bias") +PROBE_HEAD_FILE = "probe_head.safetensors" +PROBE_HEAD_KEYS = ("weight", "bias") # --------------------------------------------------------------------------- # Base weight loading @@ -832,46 +832,46 @@ def validate_untied_lm_head_saved(output_path, expected_vocab_size, hidden_size) # --------------------------------------------------------------------------- -# Classifier head transfer +# Probe head transfer # --------------------------------------------------------------------------- -def transfer_classifier_weights( - classifier_paths: list[str], +def transfer_probe_weights( + probe_paths: list[str], slot_indices: list[int], model, num_labels_per_slot: list[int], hidden_size: int, ) -> None: - """Load trained classifier heads and transfer them into the switch model. + """Load trained probe heads and transfer them into the switch model. - A classifier slot is the classifier-head analog of a LoRA adapter: it shares + A probe slot is the probe-head analog of a LoRA adapter: it shares the same control-token / adapter-index space, so its head lives in bank row - ``slot = adapter_index - 1`` of ``model.model.classifier_head``. Each - classifier directory holds a ``classifier_head.safetensors`` with a + ``slot = adapter_index - 1`` of ``model.model.probe_head``. Each + probe directory holds a ``probe_head.safetensors`` with a ``weight`` ``[n_i, hidden]`` and a ``bias`` ``[n_i]`` where ``n_i`` is this - slot's own label count. The head bank is padded to ``max_classifier_labels``. + slot's own label count. The head bank is padded to ``max_probe_labels``. Args: - classifier_paths: Directories of the classifier slots (classifier-only). + probe_paths: Directories of the probe slots (probe-only). slot_indices: Global bank row for each path (``adapter_index - 1``), - aligned with *classifier_paths*. Classifiers are the non-zero-based + aligned with *probe_paths*. Probes are the non-zero-based ``k+1..n`` suffix, so these are not ``range(len)`` and must be passed. - model: GraniteSwitch model whose ``classifier_head`` is filled in place. + model: GraniteSwitch model whose ``probe_head`` is filled in place. num_labels_per_slot: Full-width per-slot label counts - (``config.classifier_num_labels_per_slot``, indexed by + (``config.probe_num_labels_per_slot``, indexed by ``adapter_index - 1``). hidden_size: Expected hidden size (``config.hidden_size``). """ from safetensors.torch import load_file - head = getattr(model.model, "classifier_head", None) + head = getattr(model.model, "probe_head", None) if head is None: raise ValueError( - "transfer_classifier_weights was given classifier paths " - f"{classifier_paths} but the model built no classifier_head. This " + "transfer_probe_weights was given probe paths " + f"{probe_paths} but the model built no probe_head. This " "means adapter_kinds never reached the config (no slot is " - "kind=classifier). Pass adapter_kinds with a 'classifier' entry so " + "kind=probe). Pass adapter_kinds with a 'probe' entry so " "the head is constructed before loading trained weights." ) @@ -879,37 +879,37 @@ def transfer_classifier_weights( # silently corrupt another slot's head. if len(slot_indices) != len(set(slot_indices)): raise ValueError( - f"Duplicate classifier slot indices {slot_indices}; each classifier " + f"Duplicate probe slot indices {slot_indices}; each probe " "must map to a distinct bank row." ) num_rows = head.weight.shape[0] for slot in slot_indices: if not (0 <= slot < num_rows): raise ValueError( - f"Classifier slot {slot} is out of range for a classifier bank " + f"Probe slot {slot} is out of range for a probe bank " f"of {num_rows} rows (num_adapters)." ) - print(f"Loading {len(classifier_paths)} trained classifier head(s)...") + print(f"Loading {len(probe_paths)} trained probe head(s)...") with torch.no_grad(): - for path, slot in zip(classifier_paths, slot_indices): + for path, slot in zip(probe_paths, slot_indices): n_labels = num_labels_per_slot[slot] - head_file = Path(path) / CLASSIFIER_HEAD_FILE + head_file = Path(path) / PROBE_HEAD_FILE if not head_file.exists(): raise FileNotFoundError( - f"Classifier slot '{path}' is missing its trained head " - f"'{CLASSIFIER_HEAD_FILE}'. A classifier directory must " - f"contain {CLASSIFIER_HEAD_FILE} with keys " - f"{CLASSIFIER_HEAD_KEYS} (weight [n_labels, hidden], " + f"Probe slot '{path}' is missing its trained head " + f"'{PROBE_HEAD_FILE}'. A probe directory must " + f"contain {PROBE_HEAD_FILE} with keys " + f"{PROBE_HEAD_KEYS} (weight [n_labels, hidden], " f"bias [n_labels]); it has no adapter_model.safetensors." ) sd = load_file(str(head_file)) keys = set(sd.keys()) - if keys != set(CLASSIFIER_HEAD_KEYS): + if keys != set(PROBE_HEAD_KEYS): raise ValueError( f"{head_file} must contain exactly keys " - f"{set(CLASSIFIER_HEAD_KEYS)}, got {keys}." + f"{set(PROBE_HEAD_KEYS)}, got {keys}." ) weight, bias = sd["weight"], sd["bias"] @@ -918,7 +918,7 @@ def transfer_classifier_weights( f"{head_file}: weight shape {tuple(weight.shape)} does not " f"match this slot's expected (n_labels, hidden) = " f"({n_labels}, {hidden_size}). Check the head's label count " - f"against classifier_num_labels_per_slot[{slot}]={n_labels}." + f"against probe_num_labels_per_slot[{slot}]={n_labels}." ) if tuple(bias.shape) != (n_labels,): raise ValueError( @@ -936,6 +936,6 @@ def transfer_classifier_weights( bias.to(dtype=dst_w.dtype, device=dst_w.device) ) - print(f" Classifier: {path} -> classifier_head[{slot}, :{n_labels}]") + print(f" Probe: {path} -> probe_head[{slot}, :{n_labels}]") - print(f"Loaded {len(classifier_paths)} classifier head(s)") + print(f"Loaded {len(probe_paths)} probe head(s)") diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index 6de3b2ca..f5f358cb 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -28,11 +28,11 @@ class GraniteSwitchConfig(GraniteMoeHybridConfig): Args: num_adapters (int): Number of adapter slots available (LoRA and - classifier). Default: 0. Index 0 always means "base / no adapter". + probe). Default: 0. Index 0 always means "base / no adapter". adapter_token_ids (List[int]): Token IDs for adapter control. Length: num_adapters (or num_adapters + 1 with a leading base-reset token). Must be unique. adapter_token_ids[i] activates slot i+1 - (1-indexed output), a LoRA adapter or a classifier head. + (1-indexed output), a LoRA adapter or a probe head. Output 0 = base. adapter_substitute_token_ids (List[int]): Token IDs whose embeddings replace the control-token embeddings before the decoder runs. @@ -182,10 +182,10 @@ def __init__( # Shadow Residual (SR) parameters cross_stream_rank: int | None = None, dual_stream: bool = False, - # Classifier substitution parameters + # Probe substitution parameters adapter_kinds: list[str] | None = None, - classifier_label_token_ids: list[list[int] | None] | None = None, - classifier_read_layers: list[int | None] | None = None, + probe_label_token_ids: list[list[int] | None] | None = None, + probe_read_layers: list[int | None] | None = None, # vLLM residual-norm convention (for bit-exact skinning equivalence) fused_add_norm: bool = False, # Parent class defaults (Granite 4 dense configuration) @@ -387,10 +387,10 @@ def __init__( self.adapter_ranks = adapter_ranks # Each adapter slot is either a LoRA adapter ("lora", the default) or - # a classifier head ("classifier"). A classifier slot still uses a + # a probe head ("probe"). A probe slot still uses a # control token and adapter index; the switch's adapter_kind_lut routes - # its positions out of the LoRA stream and into the classifier head. - _VALID_KINDS = ("lora", "classifier") + # its positions out of the LoRA stream and into the probe head. + _VALID_KINDS = ("lora", "probe") if adapter_kinds is None: adapter_kinds = ["lora"] * num_adapters else: @@ -406,60 +406,55 @@ def __init__( f"got invalid entries {invalid}" ) self.adapter_kinds = adapter_kinds - # Classifier metadata resolved at compose time. - if classifier_label_token_ids is not None: - if len(classifier_label_token_ids) != num_adapters: + # Probe metadata resolved at compose time. + if probe_label_token_ids is not None: + if len(probe_label_token_ids) != num_adapters: raise ValueError( - f"classifier_label_token_ids length " - f"({len(classifier_label_token_ids)}) must equal num_adapters " + f"probe_label_token_ids length " + f"({len(probe_label_token_ids)}) must equal num_adapters " f"({num_adapters}); it is indexed by adapter_index - 1, with " - f"None on non-classifier slots." + f"None on non-probe slots." ) - for i, (kind, ids) in enumerate( - zip(adapter_kinds, classifier_label_token_ids) - ): - if kind == "classifier": + for i, (kind, ids) in enumerate(zip(adapter_kinds, probe_label_token_ids)): + if kind == "probe": if not ids: raise ValueError( - f"classifier slot {i} must declare its label token ids; " + f"probe slot {i} must declare its label token ids; " f"got {ids!r}." ) if any(tid < 0 for tid in ids): raise ValueError( - f"classifier_label_token_ids[{i}] must all be >= 0 " + f"probe_label_token_ids[{i}] must all be >= 0 " f"(real token ids); got {ids}" ) elif ids: raise ValueError( - f"non-classifier slot {i} must have no label token ids; " - f"got {ids!r}." + f"non-probe slot {i} must have no label token ids; got {ids!r}." ) - elif any(k == "classifier" for k in adapter_kinds): + elif any(k == "probe" for k in adapter_kinds): raise ValueError( - "adapter_kinds declares a classifier slot but " - "classifier_label_token_ids is unset. The classifier exit emits its " + "adapter_kinds declares a probe slot but " + "probe_label_token_ids is unset. The probe exit emits its " "verdict as a label word, so the per-label token ids are required. " - "Re-compose so config.json carries classifier_label_token_ids." + "Re-compose so config.json carries probe_label_token_ids." ) - self.classifier_label_token_ids = classifier_label_token_ids + self.probe_label_token_ids = probe_label_token_ids - # Per slot: zero-based decoder layer a classifier reads (switch cache + # Per slot: zero-based decoder layer a probe reads (switch cache # layers excluded); None reads the final post-norm hidden state. - if classifier_read_layers is not None: - if len(classifier_read_layers) != num_adapters: + if probe_read_layers is not None: + if len(probe_read_layers) != num_adapters: raise ValueError( - f"classifier_read_layers length ({len(classifier_read_layers)}) " + f"probe_read_layers length ({len(probe_read_layers)}) " f"must equal num_adapters ({num_adapters})." ) num_decoder_layers = self.num_hidden_layers - SWITCH_CACHE_LAYERS - for i, (kind, layer) in enumerate( - zip(adapter_kinds, classifier_read_layers) - ): + for i, (kind, layer) in enumerate(zip(adapter_kinds, probe_read_layers)): if layer is None: continue - if kind != "classifier": + if kind != "probe": raise ValueError( - f"non-classifier slot {i} must have no read layer; got {layer!r}." + f"non-probe slot {i} must have no read layer; got {layer!r}." ) # bool is an int subclass. if ( @@ -468,26 +463,26 @@ def __init__( or not 0 <= layer < num_decoder_layers ): raise ValueError( - f"classifier_read_layers[{i}] = {layer!r} must be an int in " + f"probe_read_layers[{i}] = {layer!r} must be an int in " f"0..{num_decoder_layers - 1}; omit it for the final layer." ) - self.classifier_read_layers = classifier_read_layers + self.probe_read_layers = probe_read_layers # Derived: per-slot label counts (0 on LoRA slots) and the padded bank - # width the classifier head is built with. - self.classifier_num_labels_per_slot = [ - len(ids) if ids else 0 for ids in (classifier_label_token_ids or []) + # width the probe head is built with. + self.probe_num_labels_per_slot = [ + len(ids) if ids else 0 for ids in (probe_label_token_ids or []) ] or [0] * num_adapters - self.max_classifier_labels = max(self.classifier_num_labels_per_slot, default=0) + self.max_probe_labels = max(self.probe_num_labels_per_slot, default=0) - # Derived: the control-token ids that mark a classifier slot. Both backends + # Derived: the control-token ids that mark a probe slot. Both backends # locate a verdict's read point by matching these in the original input_ids, # so the id set is defined once here. - classifier_token_offset = len(adapter_token_ids or []) - num_adapters - self.classifier_control_token_ids = [ - (adapter_token_ids or [])[i + classifier_token_offset] + probe_token_offset = len(adapter_token_ids or []) - num_adapters + self.probe_control_token_ids = [ + (adapter_token_ids or [])[i + probe_token_offset] for i, kind in enumerate(adapter_kinds) - if kind == "classifier" and adapter_token_ids is not None + if kind == "probe" and adapter_token_ids is not None ] # Default LoRA target module groups. diff --git a/src/granite_switch/hf/core/__init__.py b/src/granite_switch/hf/core/__init__.py index 28e7e412..da1a92ac 100644 --- a/src/granite_switch/hf/core/__init__.py +++ b/src/granite_switch/hf/core/__init__.py @@ -1,16 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 """Core LoRA primitives for Granite Switch (HuggingFace).""" -from .classifier import SwitchedClassifierHead from .lora import ( GraniteLoRAEmbeddedAttention, MergedSwitchedLoRALinear, SwitchedLoRALinear, ) +from .probe import SwitchedProbeHead __all__ = [ "GraniteLoRAEmbeddedAttention", "MergedSwitchedLoRALinear", - "SwitchedClassifierHead", "SwitchedLoRALinear", + "SwitchedProbeHead", ] diff --git a/src/granite_switch/hf/core/classifier.py b/src/granite_switch/hf/core/probe.py similarity index 69% rename from src/granite_switch/hf/core/classifier.py rename to src/granite_switch/hf/core/probe.py index c23d97b6..735d4218 100644 --- a/src/granite_switch/hf/core/classifier.py +++ b/src/granite_switch/hf/core/probe.py @@ -1,14 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -"""Switched classifier heads for Granite Switch. +"""Switched activation probe heads for Granite Switch. -A classifier head is an optional substitute for a LoRA adapter: an adapter slot -is either a LoRA adapter or a classifier head, selected by the same control -token and adapter index. Where a classifier is selected, the LoRA path no-ops. +An activation probe head, a linear head over a hidden state, is an optional +substitute for a LoRA adapter: an adapter slot is either a LoRA adapter or a +probe head, selected by the same control token and adapter index. Where a probe +is selected, the LoRA path no-ops. The model supplies the configured layer's hidden states (the final normalized state by default), and the head produces per-label scores. -Classifier positions are expected to be sparse, and the weight bank is replicated -on every rank at this point; revisit if classifiers ever cover most tokens or +Probe positions are expected to be sparse, and the weight bank is replicated +on every rank at this point; revisit if probes ever cover most tokens or num_labels grows large. """ @@ -16,18 +17,18 @@ import torch.nn as nn -class SwitchedClassifierHead(nn.Module): - """Per-token classifier heads selected by classifier indices. +class SwitchedProbeHead(nn.Module): + """Per-token probe heads selected by probe indices. Stacked weight bank indexed by ``(index - 1)``, like ``SwitchedLoRALinear``. - Index ``0`` = no classifier (base); ``1+`` selects that slot's head. + Index ``0`` = no probe (base); ``1+`` selects that slot's head. Positions not selected are left at zero. Args: hidden_size: Input feature dimension (the model hidden size). - num_classifier_slots: ``num_adapters`` (the full index space), so + num_probe_slots: ``num_adapters`` (the full index space), so ``index - 1`` addresses the bank directly like the LoRA bank. Slots - that are not classifiers simply never get selected. + that are not probes simply never get selected. max_num_labels: Padded label width of the bank — ``max`` over the per-slot label counts. Slots with fewer labels leave their extra rows zero. The verdict exit slices each slot back to its real label @@ -37,12 +38,12 @@ class SwitchedClassifierHead(nn.Module): def __init__( self, hidden_size: int, - num_classifier_slots: int, + num_probe_slots: int, max_num_labels: int, ): super().__init__() self.hidden_size = hidden_size - self.num_classifier_slots = num_classifier_slots + self.num_probe_slots = num_probe_slots self.num_labels = max_num_labels # Slot s (0-based) serves adapter index s + 1, matching the LoRA bank's @@ -52,21 +53,21 @@ def __init__( # head at compose time. Zero value for both an unloaded slot's row and a # loaded slot's pad rows (label counts below max_num_labels). self.weight = nn.Parameter( - torch.zeros(num_classifier_slots, max_num_labels, hidden_size) + torch.zeros(num_probe_slots, max_num_labels, hidden_size) ) - self.bias = nn.Parameter(torch.zeros(num_classifier_slots, max_num_labels)) + self.bias = nn.Parameter(torch.zeros(num_probe_slots, max_num_labels)) def forward( self, x: torch.Tensor, - classifier_indices: torch.Tensor, + probe_indices: torch.Tensor, ) -> torch.Tensor: - """Compute per-token classifier logits for selected positions. + """Compute per-token probe logits for selected positions. Args: x: Input tensor [batch, seq_len, hidden_size] or [num_tokens, hidden_size]. - classifier_indices: Per-token classifier selection [batch, seq_len] - or [num_tokens]; 0 = no classifier, 1+ = slot index. + probe_indices: Per-token probe selection [batch, seq_len] + or [num_tokens]; 0 = no probe, 1+ = slot index. Returns: Logits with the selected slot's head applied per token: @@ -83,7 +84,7 @@ def forward( logits_flat = x_flat.new_zeros((x_flat.shape[0], self.num_labels)) - idx_flat = classifier_indices.reshape(-1) + idx_flat = probe_indices.reshape(-1) # Per-token gather/scatter loop (same as SwitchedLoRALinear.forward) mask = idx_flat > 0 @@ -101,7 +102,7 @@ def forward( b = self.bias[slot] # [num_labels] x_sel = x_flat[token_indices] # [num_sel, hidden_size] - # Write (not add): a classifier replaces the LoRA. + # Write (not add): a probe replaces the LoRA. logits_flat[token_indices] = x_sel @ w.t() + b # Reshape back if needed. diff --git a/src/granite_switch/hf/modeling_granite_switch.py b/src/granite_switch/hf/modeling_granite_switch.py index 8f782ced..7a0608b1 100644 --- a/src/granite_switch/hf/modeling_granite_switch.py +++ b/src/granite_switch/hf/modeling_granite_switch.py @@ -29,12 +29,12 @@ from granite_switch.config import GraniteSwitchConfig -from .core.classifier import SwitchedClassifierHead from .core.lora import ( GraniteLoRAEmbeddedAttention, SwitchedLoRALinear, replace_shared_mlp_projections_with_lora, ) +from .core.probe import SwitchedProbeHead from .switch import create_switch logger = logging.get_logger(__name__) @@ -419,21 +419,21 @@ def __init__(self, config: GraniteSwitchConfig): self.switch = None self.adapter_token_ids = None - # Classifier heads: optional alternative for LoRA adapters. A classifier + # Probe heads: optional alternative for LoRA adapters. A probe # slot shares the adapter index/token machinery, but its index is tagged - # as "classifier" (via the switch's adapter_kind_lut). The switch's + # as "probe" (via the switch's adapter_kind_lut). The switch's # split_indices() routes those positions out of the LoRA stream (so LoRA - # no-ops there) and into a classifier stream. The head reads its + # no-ops there) and into a probe stream. The head reads its # configured layer (the final post-norm state by default), scoring # intermediate layers as the decoder runs and keeping only marker rows. - if "classifier" in getattr(config, "adapter_kinds", []): - self.classifier_head = SwitchedClassifierHead( + if "probe" in getattr(config, "adapter_kinds", []): + self.probe_head = SwitchedProbeHead( hidden_size=config.hidden_size, - num_classifier_slots=config.num_adapters, - max_num_labels=config.max_classifier_labels, + num_probe_slots=config.num_adapters, + max_num_labels=config.max_probe_labels, ) else: - self.classifier_head = None + self.probe_head = None # Decoder layers if config.num_adapters > 0: @@ -479,24 +479,22 @@ def __init__(self, config: GraniteSwitchConfig): for layer in self.layers: nn.init.zeros_(layer.cross_stream.base_layer.weight) - def _locate_classifier_markers(self, input_ids, classifier_indices, attention_mask): + def _locate_probe_markers(self, input_ids, probe_indices, attention_mask): """``(rows, pos, slots, markers_by_layer)`` for this forward's markers. ``markers_by_layer`` maps each read layer (None = final) to the marker - indices that read it. ``None`` when there is no classifier marker. Each + indices that read it. ``None`` when there is no probe marker. Each marker must be its row's last real token, since the verdict replaces that position's logits. """ - if self.classifier_head is None: + if self.probe_head is None: return None if input_ids is None: - if bool((classifier_indices > 0).any()): - raise RuntimeError( - "Classifier requests need input_ids, not inputs_embeds." - ) + if bool((probe_indices > 0).any()): + raise RuntimeError("Probe requests need input_ids, not inputs_embeds.") return None marker_ids = torch.tensor( - self.config.classifier_control_token_ids, + self.config.probe_control_token_ids, device=input_ids.device, dtype=input_ids.dtype, ) @@ -515,11 +513,11 @@ def _locate_classifier_markers(self, input_ids, classifier_indices, attention_ma if bool(misplaced.any()): bad = sorted(set(rows[misplaced].tolist())) raise RuntimeError( - f"Classifier request(s) {bad}: control token is not their last token." + f"Probe request(s) {bad}: control token is not their last token." ) - slots = classifier_indices[rows, pos] - read_layers = self.config.classifier_read_layers or [None] * len( + slots = probe_indices[rows, pos] + read_layers = self.config.probe_read_layers or [None] * len( self.config.adapter_kinds ) by_layer: dict[int | None, list[int]] = {} @@ -659,15 +657,13 @@ def forward( ) # Split the single adapter_indices stream into a LoRA stream and - # a classifier stream. Classifier positions are zeroed in lora_indices + # a probe stream. Probe positions are zeroed in lora_indices # (so the LoRA path no-ops there). if self.switch is not None: - lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices - ) + lora_indices, probe_indices = self.switch.split_indices(adapter_indices) else: lora_indices = adapter_indices - classifier_indices = torch.zeros_like(adapter_indices) + probe_indices = torch.zeros_like(adapter_indices) # Embed once, on the (possibly-rewritten) input_ids. The decoder is # token-exchange-agnostic — it just embeds whatever the switch @@ -677,12 +673,12 @@ def forward( inputs_embeds = inputs_embeds * self.embedding_multiplier # Expose adapter_indices for tests and debugging. _last_adapter_indices is - # the full stream (LoRA + classifier); _last_lora_indices is what the - # decoder routes on; _last_classifier_indices is the classifier-only + # the full stream (LoRA + probe); _last_lora_indices is what the + # decoder routes on; _last_probe_indices is the probe-only # stream the verdict reads. self._last_adapter_indices = adapter_indices self._last_lora_indices = lora_indices - self._last_classifier_indices = classifier_indices + self._last_probe_indices = probe_indices position_embeddings = None if self.rotary_emb is not None: @@ -702,18 +698,14 @@ def forward( all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None - # Resolve classifier read positions before the decoder runs, so each read + # Resolve probe read positions before the decoder runs, so each read # layer scores only its marker rows as it produces them. - markers = self._locate_classifier_markers( - input_ids, classifier_indices, attention_mask - ) - self._last_classifier_verdict = None + markers = self._locate_probe_markers(input_ids, probe_indices, attention_mask) + self._last_probe_verdict = None if markers is not None: - classifier_rows, classifier_pos, classifier_slots, markers_by_layer = ( - markers - ) - classifier_verdict = inputs_embeds.new_zeros( - (classifier_rows.numel(), self.config.max_classifier_labels) + probe_rows, probe_pos, probe_slots, markers_by_layer = markers + probe_verdict = inputs_embeds.new_zeros( + (probe_rows.numel(), self.config.max_probe_labels) ) def score_markers(layer, hidden): @@ -721,9 +713,9 @@ def score_markers(layer, hidden): if markers is None or layer not in markers_by_layer: return idx = markers_by_layer[layer] - classifier_verdict[idx] = self.classifier_head( - hidden[classifier_rows[idx], classifier_pos[idx]], - classifier_slots[idx], + probe_verdict[idx] = self.probe_head( + hidden[probe_rows[idx], probe_pos[idx]], + probe_slots[idx], ) for layer_idx, decoder_layer in enumerate(self.layers): @@ -761,11 +753,11 @@ def score_markers(layer, hidden): if markers is not None: score_markers(None, hidden_states) - self._last_classifier_verdict = ( - classifier_rows, - classifier_pos, - classifier_verdict, - classifier_slots, + self._last_probe_verdict = ( + probe_rows, + probe_pos, + probe_verdict, + probe_slots, seq_length, ) @@ -887,9 +879,9 @@ def forward( logits = self.lm_head(hidden_states[:, slice_indices, :]) logits = logits / self.config.logits_scaling - # Classifier exit: a classifier request reports its verdict as a + # Probe exit: a probe request reports its verdict as a # generated label word (like a guardian LoRA). - logits = self._apply_classifier_verdict(logits) + logits = self._apply_probe_verdict(logits) loss = None if labels is not None: @@ -908,13 +900,13 @@ def forward( attentions=outputs.attentions, ) - def _apply_classifier_verdict(self, logits: torch.Tensor) -> torch.Tensor: - """Rewrite each classifier marker's logit row to emit its label word. + def _apply_probe_verdict(self, logits: torch.Tensor) -> torch.Tensor: + """Rewrite each probe marker's logit row to emit its label word. The row is blanked to ``-inf`` and the slot's per-label verdict scores are - placed on its label token ids (``classifier_label_token_ids``). + placed on its label token ids (``probe_label_token_ids``). """ - pending = getattr(self.model, "_last_classifier_verdict", None) + pending = getattr(self.model, "_last_probe_verdict", None) if pending is None: return logits rows, pos, verdict, slots, seq_length = pending @@ -923,12 +915,12 @@ def _apply_classifier_verdict(self, logits: torch.Tensor) -> torch.Tensor: cols = pos - (seq_length - logits.shape[1]) if bool((cols < 0).any()): raise RuntimeError( - "A classifier marker's logits were not kept (logits_to_keep)." + "A probe marker's logits were not kept (logits_to_keep)." ) # Each marker writes only its own slot's label ids and only that slot's # real label count (verdict[i, :n]); padded columns are never read. - per_slot_label_ids = self.config.classifier_label_token_ids + per_slot_label_ids = self.config.probe_label_token_ids logits[rows, cols, :] = float("-inf") for i, (r, c, slot) in enumerate( zip(rows.tolist(), cols.tolist(), slots.tolist()) diff --git a/src/granite_switch/hf/switch/multi.py b/src/granite_switch/hf/switch/multi.py index 3f8f48ae..0ec61111 100644 --- a/src/granite_switch/hf/switch/multi.py +++ b/src/granite_switch/hf/switch/multi.py @@ -66,7 +66,7 @@ from transformers.cache_utils import Cache from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from ...classifier_routing import build_adapter_kind_lut, split_adapter_indices +from ...probe_routing import build_adapter_kind_lut, split_adapter_indices from ...token_exchange import apply_token_exchange, build_control_to_substitute_lut from .codes import KerdockDGCodeGenerator, recover_count_from_signal @@ -276,9 +276,9 @@ def __init__( # non-persistent buffer is zeroed by ``from_pretrained``. This LUT uses # -1 as the "not a control token" sentinel, so an all-zero LUT would # rewrite EVERY token id to 0 in apply_token_exchange. - # adapter_kind_lut: marks which adapter indices are classifier slots, so + # adapter_kind_lut: marks which adapter indices are probe slots, so # split_indices can route those positions out of the LoRA stream. None when - # no slot is a classifier. persistent=True so it survives HF's + # no slot is a probe. persistent=True so it survives HF's # meta-device loading, like the token-exchange LUT below. kind_lut = build_adapter_kind_lut(config) if kind_lut is not None: @@ -322,9 +322,9 @@ def _apply(self, *args, **kwargs): def split_indices( self, adapter_indices: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """Split ``[batch, seq_len]`` or ``[num_tokens]`` into (lora, classifier). + """Split ``[batch, seq_len]`` or ``[num_tokens]`` into (lora, probe). - See :func:`granite_switch.classifier_routing.split_adapter_indices`. + See :func:`granite_switch.probe_routing.split_adapter_indices`. """ return split_adapter_indices(self.adapter_kind_lut, adapter_indices) diff --git a/src/granite_switch/classifier_routing.py b/src/granite_switch/probe_routing.py similarity index 50% rename from src/granite_switch/classifier_routing.py rename to src/granite_switch/probe_routing.py index 5467291c..0c4c23f3 100644 --- a/src/granite_switch/classifier_routing.py +++ b/src/granite_switch/probe_routing.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -"""Classifier-slot routing helpers, shared by both backends' switches. +"""Activation-probe slot routing helpers, shared by both backends' switches. -An adapter slot is either a generation adapter or a classifier head, selected -by the same control token and adapter index. These helpers split the switch's -single per-token index stream so that classifier positions no-op the LoRA path -and feed the classifier head instead. +An adapter slot is either a generation adapter or an activation probe head, +selected by the same control token and adapter index. These helpers split the +switch's single per-token index stream so that probe positions no-op the LoRA +path and feed the probe head instead. Backend-neutral on purpose: the table is a pure function of ``config.adapter_kinds``, so both switch engines (HF and vLLM MultiSwitch) @@ -15,20 +15,20 @@ def build_adapter_kind_lut(config) -> torch.Tensor | None: - """Bool LUT over adapter indices marking which slots are classifiers. + """Bool LUT over adapter indices marking which slots are probes. ``[num_adapters + 1]``, indexed by ADAPTER INDEX: entry ``i`` is True when - adapter index ``i`` is a classifier slot (index 0, base, is always False). - ``None`` when no slot is a classifier, which lets the switch skip the split + adapter index ``i`` is a probe slot (index 0, base, is always False). + ``None`` when no slot is a probe, which lets the switch skip the split entirely. Built from ``config.adapter_kinds``, so both backends' switches agree on the table by construction. """ kinds = getattr(config, "adapter_kinds", None) if config is not None else None - if kinds is None or not any(k == "classifier" for k in kinds): + if kinds is None or not any(k == "probe" for k in kinds): return None lut = torch.zeros(len(kinds) + 1, dtype=torch.bool) for slot, kind in enumerate(kinds): # slot is the 0-based adapter position - if kind == "classifier": + if kind == "probe": lut[slot + 1] = True # adapter index is slot + 1 return lut @@ -36,23 +36,23 @@ def build_adapter_kind_lut(config) -> torch.Tensor | None: def split_adapter_indices( kind_lut: torch.Tensor | None, adapter_indices: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """Split a per-token index stream into LoRA and classifier streams. + """Split a per-token index stream into LoRA and probe streams. Shapes are preserved (``[batch, seq_len]`` on HF, ``[total_tokens]`` on vLLM), so this serves both backends: - - ``lora_indices``: classifier positions zeroed, so the LoRA path no-ops + - ``lora_indices``: probe positions zeroed, so the LoRA path no-ops there. - - ``classifier_indices``: the adapter index at classifier positions, else 0. + - ``probe_indices``: the adapter index at probe positions, else 0. - With ``kind_lut`` None (no classifier slots) the LoRA stream passes through - and the classifier stream is all zeros. + With ``kind_lut`` None (no probe slots) the LoRA stream passes through + and the probe stream is all zeros. """ if kind_lut is None: return adapter_indices, torch.zeros_like(adapter_indices) - is_classifier = kind_lut[adapter_indices] + is_probe = kind_lut[adapter_indices] zeros = torch.zeros_like(adapter_indices) return ( - torch.where(is_classifier, zeros, adapter_indices), - torch.where(is_classifier, adapter_indices, zeros), + torch.where(is_probe, zeros, adapter_indices), + torch.where(is_probe, adapter_indices, zeros), ) diff --git a/src/granite_switch/vllm/core/__init__.py b/src/granite_switch/vllm/core/__init__.py index 894e7eb6..eefeb042 100644 --- a/src/granite_switch/vllm/core/__init__.py +++ b/src/granite_switch/vllm/core/__init__.py @@ -8,13 +8,13 @@ - lora_kernel_meta: Bitmask + remap metadata for the fused kernel """ -from .classifier import SwitchedClassifierHead from .lora import SwitchedLoRALinear from .lora_kernel_meta import FusedLoRAKernelMeta, LoRAContext +from .probe import SwitchedProbeHead __all__ = [ "FusedLoRAKernelMeta", "LoRAContext", - "SwitchedClassifierHead", "SwitchedLoRALinear", + "SwitchedProbeHead", ] diff --git a/src/granite_switch/vllm/core/classifier.py b/src/granite_switch/vllm/core/probe.py similarity index 62% rename from src/granite_switch/vllm/core/classifier.py rename to src/granite_switch/vllm/core/probe.py index c3e018ec..140b332d 100644 --- a/src/granite_switch/vllm/core/classifier.py +++ b/src/granite_switch/vllm/core/probe.py @@ -1,14 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -"""Switched classifier heads for Granite Switch (vLLM backend). +"""Switched activation probe heads for Granite Switch (vLLM backend). -A classifier head is an optional substitute for a LoRA adapter: an adapter slot -is either a LoRA adapter or a classifier head, selected by the same control -token and adapter index. Where a classifier is selected, the LoRA path no-ops. +An activation probe head, a linear head over a hidden state, is an optional +substitute for a LoRA adapter: an adapter slot is either a LoRA adapter or a +probe head, selected by the same control token and adapter index. Where a probe +is selected, the LoRA path no-ops. The model supplies the configured layer's hidden states (the final normalized state by default), and the head produces per-label scores. -Classifier positions are expected to be sparse, and the weight bank is replicated -on every rank at this point; revisit if classifiers ever cover most tokens or +Probe positions are expected to be sparse, and the weight bank is replicated +on every rank at this point; revisit if probes ever cover most tokens or num_labels grows large. """ @@ -16,11 +17,11 @@ import torch.nn as nn -class SwitchedClassifierHead(nn.Module): - """Per-token classifier heads selected by classifier indices (vLLM). +class SwitchedProbeHead(nn.Module): + """Per-token probe heads selected by probe indices (vLLM). Compile-safe, on the flat ``[total_tokens, hidden_size]`` layout vLLM's - scheduler produces. Index ``0`` = no classifier (base -> zero logits); + scheduler produces. Index ``0`` = no probe (base -> zero logits); ``1+`` selects that slot's head. The weight bank is stacked ``[num_slots, ...]`` and indexed by @@ -28,9 +29,9 @@ class SwitchedClassifierHead(nn.Module): Args: hidden_size: Input feature dimension (the model hidden size). - num_classifier_slots: ``num_adapters`` (the full index space). Slot ``s`` + num_probe_slots: ``num_adapters`` (the full index space). Slot ``s`` (0-based) serves adapter index ``s + 1``. Slots that are not - classifiers simply never get selected. + probes simply never get selected. max_num_labels: Padded label width of the bank (``max`` over the per-slot label counts). Keeping this uniform is what lets the per-token logits be one fused ``einsum``. Slots with fewer labels leave their extra @@ -41,12 +42,12 @@ class SwitchedClassifierHead(nn.Module): def __init__( self, hidden_size: int, - num_classifier_slots: int, + num_probe_slots: int, max_num_labels: int, ): super().__init__() self.hidden_size = hidden_size - self.num_classifier_slots = num_classifier_slots + self.num_probe_slots = num_probe_slots self.num_labels = max_num_labels # Stacked bank indexed by (adapter_idx - 1), like the LoRA bank @@ -57,30 +58,30 @@ def __init__( # head at compose time. Zero value for both an unloaded slot's row and a # loaded slot's pad rows (label counts below max_num_labels). self.weight = nn.Parameter( - torch.zeros(num_classifier_slots, max_num_labels, hidden_size) + torch.zeros(num_probe_slots, max_num_labels, hidden_size) ) - self.bias = nn.Parameter(torch.zeros(num_classifier_slots, max_num_labels)) + self.bias = nn.Parameter(torch.zeros(num_probe_slots, max_num_labels)) def forward( self, x: torch.Tensor, - classifier_indices: torch.Tensor, + probe_indices: torch.Tensor, ) -> torch.Tensor: - """Compute per-token classifier logits for selected positions. + """Compute per-token probe logits for selected positions. Branchless and loop-free so it traces cleanly under ``@support_torch_compile``. Args: x: Input tensor [total_tokens, hidden_size]. - classifier_indices: [total_tokens]; 0 = no classifier, 1+ = slot + probe_indices: [total_tokens]; 0 = no probe, 1+ = slot index (the un-offset adapter index). Returns: Logits [total_tokens, num_labels]. Positions with index 0 are zero. """ - is_classifier = classifier_indices > 0 # [total_tokens] - slot = (classifier_indices - 1).clamp_min(0) # [total_tokens]; base -> row 0 + is_probe = probe_indices > 0 # [total_tokens] + slot = (probe_indices - 1).clamp_min(0) # [total_tokens]; base -> row 0 w_per_token = self.weight[slot] # [total_tokens, num_labels, hidden] b_per_token = self.bias[slot] # [total_tokens, num_labels] @@ -90,5 +91,5 @@ def forward( logits = logits + b_per_token # Zero out base positions. - logits = logits * is_classifier.unsqueeze(-1).to(logits.dtype) + logits = logits * is_probe.unsqueeze(-1).to(logits.dtype) return logits diff --git a/src/granite_switch/vllm/decoder/interface.py b/src/granite_switch/vllm/decoder/interface.py index 81fe9c44..dbcc93cd 100644 --- a/src/granite_switch/vllm/decoder/interface.py +++ b/src/granite_switch/vllm/decoder/interface.py @@ -596,7 +596,7 @@ def layer_output(self, state, adapter_indices): hs = state.hs m = adapter_indices.shape[0] # Read after the layer's cross-stream shunt: adapter-active tokens use - # the adapted stream, the rest (including classifier markers, whose LoRA + # the adapted stream, the rest (including probe markers, whose LoRA # stream is zeroed) use the base stream. select = (adapter_indices > 0).unsqueeze(1) # [M, 1] return torch.where(select, hs[m:], hs[:m]) diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index eeba6c81..853aa0f6 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -168,68 +168,64 @@ def __init__( self.lora_meta = None self.lora_ctx = None - # Classifier heads: optional alternative for LoRA adapters. A classifier + # Probe heads: optional alternative for LoRA adapters. A probe # slot shares the adapter index/token machinery, but its index is tagged - # as "classifier" (via the switch's adapter_kind_lut). The switch's + # as "probe" (via the switch's adapter_kind_lut). The switch's # split_indices() routes those positions out of the LoRA stream (so LoRA - # no-ops there) and into a classifier stream. The head scores each slot + # no-ops there) and into a probe stream. The head scores each slot # at its read layer (the final post-norm state by default). - if "classifier" in getattr(config, "adapter_kinds", []): - from .core import SwitchedClassifierHead + if "probe" in getattr(config, "adapter_kinds", []): + from .core import SwitchedProbeHead - self.classifier_head = SwitchedClassifierHead( + self.probe_head = SwitchedProbeHead( hidden_size=config.hidden_size, - num_classifier_slots=config.num_adapters, - max_num_labels=config.max_classifier_labels, + num_probe_slots=config.num_adapters, + max_num_labels=config.max_probe_labels, ) # Packed verdicts assume the ordinary sampler's row selection. if vllm_config.speculative_config is not None: - raise ValueError( - "Classifier slots do not support speculative decoding." - ) + raise ValueError("Probe slots do not support speculative decoding.") if vllm_config.parallel_config.use_ubatching: - raise ValueError("Classifier slots do not support microbatching.") + raise ValueError("Probe slots do not support microbatching.") # Built once: copying a list to CUDA inside forward breaks capture. self.register_buffer( - "classifier_control_token_ids_t", - torch.tensor(config.classifier_control_token_ids, dtype=torch.long), + "probe_control_token_ids_t", + torch.tensor(config.probe_control_token_ids, dtype=torch.long), persistent=False, ) - is_classifier = [kind == "classifier" for kind in config.adapter_kinds] + is_probe = [kind == "probe" for kind in config.adapter_kinds] # Index 0 is a sentinel so a 1-based slot id indexes it directly. self.register_buffer( - "classifier_slot_mask_t", - torch.tensor([False, *is_classifier]), + "probe_slot_mask_t", + torch.tensor([False, *is_probe]), persistent=False, ) # One buffer of slot ids per read layer (None = final). Store buffer # names, not tensors: Module.to() replaces registered buffers. - read_layers = config.classifier_read_layers or [None] * num_adapters + read_layers = config.probe_read_layers or [None] * num_adapters slots_by_layer: dict[int | None, list[int]] = {} - for slot, (is_clf, layer) in enumerate( - zip(is_classifier, read_layers), start=1 - ): + for slot, (is_clf, layer) in enumerate(zip(is_probe, read_layers), start=1): if is_clf: slots_by_layer.setdefault(layer, []).append(slot) - self._classifier_final_layer_group = None - self._classifier_layer_groups: dict[int, str] = {} + self._probe_final_layer_group = None + self._probe_layer_groups: dict[int, str] = {} for layer, slot_ids in slots_by_layer.items(): - name = f"_classifier_slots_{'final' if layer is None else layer}" + name = f"_probe_slots_{'final' if layer is None else layer}" self.register_buffer( name, torch.tensor(slot_ids, dtype=torch.long), persistent=False ) if layer is None: - self._classifier_final_layer_group = name + self._probe_final_layer_group = name else: - self._classifier_layer_groups[layer] = name + self._probe_layer_groups[layer] = name else: - self.classifier_head = None - self._classifier_final_layer_group = None - self._classifier_layer_groups = {} + self.probe_head = None + self._probe_final_layer_group = None + self._probe_layer_groups = {} # 3. Base transformer layers with custom LoRA # @@ -309,12 +305,12 @@ def make_empty_intermediate_tensors( device=device, ), } - if self.classifier_head is not None: - tensors["classifier_indices"] = torch.zeros( + if self.probe_head is not None: + tensors["probe_indices"] = torch.zeros( batch_size, dtype=torch.long, device=device ) - tensors["classifier_scores"] = torch.zeros( - (batch_size, self.config.max_classifier_labels), + tensors["probe_scores"] = torch.zeros( + (batch_size, self.config.max_probe_labels), dtype=dtype, device=device, ) @@ -394,26 +390,24 @@ def forward( adapter_indices = intermediate_tensors["adapter_indices"] # Split the single adapter_indices stream into a LoRA stream and a - # classifier stream. Classifier positions are zeroed in lora_indices (so + # probe stream. Probe positions are zeroed in lora_indices (so # the LoRA path no-ops there). if self.switch is not None: - lora_indices, classifier_indices = self.switch.split_indices( - adapter_indices - ) + lora_indices, probe_indices = self.switch.split_indices(adapter_indices) else: lora_indices = adapter_indices - classifier_indices = torch.zeros_like(adapter_indices) + probe_indices = torch.zeros_like(adapter_indices) # The verdict is read only at the marker itself, not at tokens that # inherit its slot afterwards. Later PP ranks have no input ids, so they # receive the marker-only indices from the first rank. - if self.classifier_head is not None: + if self.probe_head is not None: if not get_pp_group().is_first_rank: - classifier_indices = intermediate_tensors["classifier_indices"] + probe_indices = intermediate_tensors["probe_indices"] elif input_ids is not None: - is_marker = torch.isin(input_ids, self.classifier_control_token_ids_t) - classifier_indices = torch.where( - is_marker, classifier_indices, torch.zeros_like(classifier_indices) + is_marker = torch.isin(input_ids, self.probe_control_token_ids_t) + probe_indices = torch.where( + is_marker, probe_indices, torch.zeros_like(probe_indices) ) # Rebuild this rank's fixed-size kernel metadata from token-leading @@ -449,28 +443,28 @@ def forward( # threads (hidden_states, residual); SR re-stacks to [2M,H] rank-locally # in enter_decoder_stack and collapses it in exit_decoder_stack. All # per-forward metadata (kernel meta + hiding masks) is on the shared ctx. - # Classifier slots reading an intermediate layer are scored as soon as + # Probe slots reading an intermediate layer are scored as soon as # that layer runs; only the small score tensor is carried forward. - if self.classifier_head is not None: + if self.probe_head is not None: if get_pp_group().is_first_rank: - classifier_scores = torch.zeros( - classifier_indices.shape[0], - self.config.max_classifier_labels, - device=classifier_indices.device, + probe_scores = torch.zeros( + probe_indices.shape[0], + self.config.max_probe_labels, + device=probe_indices.device, dtype=hidden_states.dtype, ) else: - classifier_scores = intermediate_tensors["classifier_scores"] + probe_scores = intermediate_tensors["probe_scores"] state = self.decoder_interface.enter_decoder_stack(hidden_states, residual) for i in range(self.start_layer, self.end_layer): state = self.decoder_interface.run_layer(self.layers[i], positions, state) - group = self._classifier_layer_groups.get(i) + group = self._probe_layer_groups.get(i) if group is not None: - classifier_scores = self._score_group( + probe_scores = self._score_group( group, self.decoder_interface.layer_output(state, lora_indices), - classifier_indices, - classifier_scores, + probe_indices, + probe_scores, ) if get_pp_group().is_last_rank: @@ -485,12 +479,12 @@ def forward( # Pack scores and slot ids onto the hidden states so vLLM's own # logits_indices gather selects them; compute_logits unpacks them. - if self.classifier_head is not None: - classifier_scores, classifier_slot = self._compute_classifier_outputs( - hidden_states, classifier_indices, classifier_scores + if self.probe_head is not None: + probe_scores, probe_slot = self._compute_probe_outputs( + hidden_states, probe_indices, probe_scores ) hidden_states = torch.cat( - [hidden_states, classifier_scores, classifier_slot.unsqueeze(-1)], + [hidden_states, probe_scores, probe_slot.unsqueeze(-1)], dim=-1, ) @@ -505,46 +499,46 @@ def forward( intermediate_tensors["hidden_states"] = h_out intermediate_tensors["residual"] = r_out intermediate_tensors["adapter_indices"] = adapter_indices - if self.classifier_head is not None: - intermediate_tensors["classifier_indices"] = classifier_indices - intermediate_tensors["classifier_scores"] = classifier_scores + if self.probe_head is not None: + intermediate_tensors["probe_indices"] = probe_indices + intermediate_tensors["probe_scores"] = probe_scores return intermediate_tensors def _score_group( self, group: str, hidden: torch.Tensor, - classifier_indices: torch.Tensor, - classifier_scores: torch.Tensor, + probe_indices: torch.Tensor, + probe_scores: torch.Tensor, ) -> torch.Tensor: """Write the scores of the slots in buffer ``group``, read from ``hidden``.""" - in_group = torch.isin(classifier_indices, getattr(self, group)) - slot_logits = self.classifier_head(hidden, classifier_indices) - return torch.where(in_group.unsqueeze(-1), slot_logits, classifier_scores) + in_group = torch.isin(probe_indices, getattr(self, group)) + slot_logits = self.probe_head(hidden, probe_indices) + return torch.where(in_group.unsqueeze(-1), slot_logits, probe_scores) - def _compute_classifier_outputs( + def _compute_probe_outputs( self, final_hidden_states: torch.Tensor, - classifier_indices: torch.Tensor, - classifier_scores: torch.Tensor, + probe_indices: torch.Tensor, + probe_scores: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Add final-layer scores and return per-token (scores, slot ids). - ``classifier_indices`` is nonzero only at this forward's marker tokens, + ``probe_indices`` is nonzero only at this forward's marker tokens, computed on the first PP rank and passed unchanged between stages. - Callers must place one classifier marker last in each request; + Callers must place one probe marker last in each request; misplaced or duplicate markers are not rejected (request-boundary validation is deferred). """ - classifier_slot = classifier_indices.to(final_hidden_states.dtype) - if self._classifier_final_layer_group is not None: - classifier_scores = self._score_group( - self._classifier_final_layer_group, + probe_slot = probe_indices.to(final_hidden_states.dtype) + if self._probe_final_layer_group is not None: + probe_scores = self._score_group( + self._probe_final_layer_group, final_hidden_states, - classifier_indices, - classifier_scores, + probe_indices, + probe_scores, ) - return classifier_scores, classifier_slot + return probe_scores, probe_slot @MULTIMODAL_REGISTRY.register_processor( @@ -729,7 +723,7 @@ def forward( intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: - """Forward pass returning hidden states (classifier columns packed on).""" + """Forward pass returning hidden states (probe columns packed on).""" return self.model( input_ids=input_ids, positions=positions, @@ -752,15 +746,15 @@ def compute_logits( have to act where adapter_indices and hidden_states still share the same token dimension. - Classifier exit: the classifier request reports its verdict as a - generated token, like a guardian LoRA. With a classifier head the + Probe exit: the probe request reports its verdict as a + generated token, like a guardian LoRA. With a probe head the hidden states carry packed (scores, slot id) columns; split them off, - run the LM head on the real hidden state, then rewrite each classifier + run the LM head on the real hidden state, then rewrite each probe row: the whole vocab to -inf and the per-label scores onto the label - words' token ids (``classifier_label_token_ids``). Other rows keep + words' token ids (``probe_label_token_ids``). Other rows keep their LM logits. """ - if self.model.classifier_head is not None: + if self.model.probe_head is not None: hidden_size = self.config.hidden_size real_hidden = hidden_states[:, :hidden_size] verdict = hidden_states[:, hidden_size:-1] @@ -769,15 +763,15 @@ def compute_logits( real_hidden = hidden_states verdict = slot = None logits = self.logits_processor(self.lm_head, real_hidden) - return self._apply_classifier_verdict(logits, verdict, slot) + return self._apply_probe_verdict(logits, verdict, slot) - def _apply_classifier_verdict( + def _apply_probe_verdict( self, logits: torch.Tensor | None, verdict: torch.Tensor | None, slot: torch.Tensor | None, ) -> torch.Tensor | None: - """Rewrite classifier requests' logit rows to emit their label word. + """Rewrite probe requests' logit rows to emit their label word. Rows are classified by the slot mask, not just ``slot > 0``: vLLM's profiling sampler run feeds random hidden states, so the packed slot @@ -786,19 +780,19 @@ def _apply_classifier_verdict( if logits is None or verdict is None or slot is None: return logits - slot_mask = self.model.classifier_slot_mask_t + slot_mask = self.model.probe_slot_mask_t num_slots = slot_mask.shape[0] - 1 slot_int = slot.round().long() # Out-of-range values land on the sentinel row 0. safe_idx = slot_int.clamp(min=0, max=num_slots) - is_classifier = (slot_int == safe_idx) & slot_mask[safe_idx] - if not bool(is_classifier.any()): + is_probe = (slot_int == safe_idx) & slot_mask[safe_idx] + if not bool(is_probe.any()): return logits # Each row writes only its own slot's real label count; padded # columns are never read. - per_slot_label_ids = self.config.classifier_label_token_ids - rows = is_classifier.nonzero(as_tuple=True)[0] + per_slot_label_ids = self.config.probe_label_token_ids + rows = is_probe.nonzero(as_tuple=True)[0] logits[rows] = float("-inf") for r in rows.tolist(): row_slot = int(slot_int[r]) diff --git a/src/granite_switch/vllm/switch/multi.py b/src/granite_switch/vllm/switch/multi.py index 7c7995f6..a2292706 100644 --- a/src/granite_switch/vllm/switch/multi.py +++ b/src/granite_switch/vllm/switch/multi.py @@ -98,14 +98,14 @@ from vllm.config import VllmConfig from vllm.model_executor.layers.attention.attention import Attention -from granite_switch.classifier_routing import ( - build_adapter_kind_lut, - split_adapter_indices, -) from granite_switch.hf.switch.codes import ( KerdockDGCodeGenerator, recover_count_from_signal, ) +from granite_switch.probe_routing import ( + build_adapter_kind_lut, + split_adapter_indices, +) from granite_switch.token_exchange import ( apply_token_exchange, build_control_to_substitute_lut, @@ -316,9 +316,9 @@ def __init__( # non-persistent buffer is zeroed by checkpoint loading. This LUT uses -1 # as the "not a control token" sentinel, so an all-zero LUT would rewrite # EVERY token id to 0 in apply_token_exchange. - # adapter_kind_lut: marks which adapter indices are classifier slots, so + # adapter_kind_lut: marks which adapter indices are probe slots, so # split_indices can route those positions out of the LoRA stream. None when - # no slot is a classifier. + # no slot is a probe. kind_lut = build_adapter_kind_lut(hf_config) if kind_lut is not None: self.register_buffer("adapter_kind_lut", kind_lut, persistent=False) @@ -334,9 +334,9 @@ def __init__( def split_indices( self, adapter_indices: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """Split ``[total_tokens]`` into (lora_indices, classifier_indices). + """Split ``[total_tokens]`` into (lora_indices, probe_indices). - See :func:`granite_switch.classifier_routing.split_adapter_indices`. + See :func:`granite_switch.probe_routing.split_adapter_indices`. """ return split_adapter_indices(self.adapter_kind_lut, adapter_indices) diff --git a/tests/composer/test_adapter_loader.py b/tests/composer/test_adapter_loader.py index e27e316f..9cb54254 100644 --- a/tests/composer/test_adapter_loader.py +++ b/tests/composer/test_adapter_loader.py @@ -537,9 +537,9 @@ def test_fallback_precedence_and_yaml_parity(self, tmp_path): yaml.dump(manifest_data, f) input_path = str(manifest_file) - yaml_adapters, classifier_meta = discover_adapters_from_yaml(input_path) - # No entry declared kind/labels, so no classifier metadata. - assert classifier_meta == {} + yaml_adapters, probe_meta = discover_adapters_from_yaml(input_path) + # No entry declared kind/labels, so no probe metadata. + assert probe_meta == {} # Verification of Parity: compare (path, name, technology) - source differs by design # discover_adapters returns source=None, discover_adapters_from_yaml returns source=manifest_path @@ -551,8 +551,8 @@ def test_fallback_precedence_and_yaml_parity(self, tmp_path): assert sorted(adapters_without_source) == sorted(yaml_without_source) -class TestManifestClassifierMetadata: - """discover_adapters_from_yaml surfaces optional classifier declarations.""" +class TestManifestProbeMetadata: + """discover_adapters_from_yaml surfaces optional probe declarations.""" def test_kind_and_labels_are_parsed(self, tmp_path): manifest = tmp_path / "manifest.yaml" @@ -562,29 +562,29 @@ def test_kind_and_labels_are_parsed(self, tmp_path): "hallucination": { "path": "/p/hall", "type": "lora", - "kind": "classifier", + "kind": "probe", "labels": ["faithful", "hallucinated"], }, }, manifest.open("w"), ) - found, classifier_meta = discover_adapters_from_yaml(str(manifest)) + found, probe_meta = discover_adapters_from_yaml(str(manifest)) # The adapter tuples are unchanged (still 4-wide). assert {(n) for _, n, _, _ in found} == {"rag", "hallucination"} - # Only the classifier entry appears in the metadata dict. - assert classifier_meta == { + # Only the probe entry appears in the metadata dict. + assert probe_meta == { "hallucination": { - "kind": "classifier", + "kind": "probe", "labels": ["faithful", "hallucinated"], "layer": None, } } - def test_plain_manifest_has_no_classifier_metadata(self, tmp_path): + def test_plain_manifest_has_no_probe_metadata(self, tmp_path): manifest = tmp_path / "manifest.yaml" yaml.dump({"rag": {"path": "/p/rag", "type": "lora"}}, manifest.open("w")) - _found, classifier_meta = discover_adapters_from_yaml(str(manifest)) - assert classifier_meta == {} + _found, probe_meta = discover_adapters_from_yaml(str(manifest)) + assert probe_meta == {} diff --git a/tests/composer/test_base_reset_token.py b/tests/composer/test_base_reset_token.py index 3869739c..8f677d6d 100644 --- a/tests/composer/test_base_reset_token.py +++ b/tests/composer/test_base_reset_token.py @@ -22,8 +22,8 @@ _PROBE = ( "granite_switch.composer.compose_granite_switch._probe_lora_substitute_token_id" ) -_CLASSIFIER_SUB = ( - "granite_switch.composer.compose_granite_switch._classifier_substitute_token_id" +_PROBE_SLOT_SUB = ( + "granite_switch.composer.compose_granite_switch._probe_slot_substitute_token_id" ) _ALORA = "granite_switch.composer.tokenizer_setup.get_alora_first_invocation_token_id" @@ -55,7 +55,7 @@ def _run(self, base_reset): tokenizer = MockTokenizer(initial_vocab_size=500) with ( patch(_PROBE, return_value=42), - patch(_CLASSIFIER_SUB, return_value=43), + patch(_PROBE_SLOT_SUB, return_value=43), patch(_ALORA, return_value=77), ): return build_control_token_lists(tokenizer, self._ADAPTERS, base_reset) @@ -93,7 +93,7 @@ def test_compose_layout_is_the_one_the_engine_reads_as_base_reset(self): tokenizer = _Tok(initial_vocab_size=500) with ( patch(_PROBE, return_value=42), - patch(_CLASSIFIER_SUB, return_value=43), + patch(_PROBE_SLOT_SUB, return_value=43), patch(_ALORA, return_value=77), ): token_ids, _special, _subs = build_control_token_lists( @@ -138,7 +138,7 @@ def test_default_layout_keeps_the_adapter_offset(self): tokenizer = _Tok(initial_vocab_size=500) with ( patch(_PROBE, return_value=42), - patch(_CLASSIFIER_SUB, return_value=43), + patch(_PROBE_SLOT_SUB, return_value=43), patch(_ALORA, return_value=77), ): token_ids, _special, _subs = build_control_token_lists( diff --git a/tests/composer/test_control_lut_refresh.py b/tests/composer/test_control_lut_refresh.py index d722969c..08561779 100644 --- a/tests/composer/test_control_lut_refresh.py +++ b/tests/composer/test_control_lut_refresh.py @@ -132,24 +132,24 @@ def test_no_mapping_is_not_an_error(switch_type): assert refresh_switch_control_lut(model) is False -def test_classifier_routing_survives_audio_resize_and_strict_reload(tmp_path): - """Classifier + audio: every retained routing table must match the saved config. +def test_probe_routing_survives_audio_resize_and_strict_reload(tmp_path): + """Probe + audio: every retained routing table must match the saved config. The control LUT is vocabulary-sized and is rebuilt after the resize; the - classifier kind LUT is indexed by adapter, so the resize cannot stale it. + probe kind LUT is indexed by adapter, so the resize cannot stale it. A strict reload (no ``ignore_mismatched_sizes``) must keep both, and the - classifier marker must still route to its head and emit a label word. + probe marker must still route to its head and emit a label word. """ from granite_switch.composer.validator import validate_control_lut config = _config( "multi", - adapter_kinds=["lora", "classifier"], - classifier_label_token_ids=[None, [10, 11]], + adapter_kinds=["lora", "probe"], + probe_label_token_ids=[None, [10, 11]], ) model = GraniteSwitchForCausalLM(config).eval() with torch.no_grad(): - model.model.classifier_head.bias[1, :2] = torch.tensor([-5.0, 5.0]) + model.model.probe_head.bias[1, :2] = torch.tensor([-5.0, 5.0]) model.resize_token_embeddings(_VOCAB_WITH_AUDIO) assert refresh_switch_control_lut(model) is True validate_control_lut(model) @@ -166,5 +166,5 @@ def test_classifier_routing_survives_audio_resize_and_strict_reload(tmp_path): with torch.no_grad(): out = loaded(input_ids=torch.tensor([[1, 2, 3, _CTRL_IDS[1]]]), use_cache=False) assert loaded.model._last_lora_indices.tolist() == [[0, 0, 0, 0]] - assert loaded.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] + assert loaded.model._last_probe_indices.tolist() == [[0, 0, 0, 2]] assert int(out.logits[0, -1].argmax()) == 11 diff --git a/tests/composer/test_multi_audio_compose_e2e.py b/tests/composer/test_multi_audio_compose_e2e.py index 339644d9..a2b1de5e 100644 --- a/tests/composer/test_multi_audio_compose_e2e.py +++ b/tests/composer/test_multi_audio_compose_e2e.py @@ -156,30 +156,30 @@ def test_multi_audio_checkpoint_loads_without_size_overrides(audio_multi): # --------------------------------------------------------------------------- # -# A LoRA + classifier checkpoint composed with and without --enable-audio must -# load strictly (no ignore_mismatched_sizes) and still route the classifier +# A LoRA + probe checkpoint composed with and without --enable-audio must +# load strictly (no ignore_mismatched_sizes) and still route the probe # marker to its head. The audio marker grows the vocabulary past the control # ids, so this also checks the saved control LUT matches the shipped config. # --------------------------------------------------------------------------- # -_CLASSIFIER_BASE = "ibm-granite/granite-4.1-3b" -_CLASSIFIER_LABELS = ["safe", "unsafe"] +_PROBE_BASE = "ibm-granite/granite-4.1-3b" +_PROBE_LABELS = ["safe", "unsafe"] -def _write_classifier_manifest(root): +def _write_probe_manifest(root): """A real LoRA plus a synthetic bias-only head (label 1 always wins).""" import torch from safetensors.torch import save_file from transformers import AutoConfig - from .test_classifier_compose import _find_rag_lora_dir + from .test_probe_compose import _find_rag_lora_dir - hidden = AutoConfig.from_pretrained(_CLASSIFIER_BASE).hidden_size + hidden = AutoConfig.from_pretrained(_PROBE_BASE).hidden_size head_dir = root / "safety" / "safety" / "granite-4.1-3b" / "lora" head_dir.mkdir(parents=True) bias = torch.tensor([-10.0, 10.0]) save_file( - {"weight": torch.zeros(len(_CLASSIFIER_LABELS), hidden), "bias": bias}, - str(head_dir / "classifier_head.safetensors"), + {"weight": torch.zeros(len(_PROBE_LABELS), hidden), "bias": bias}, + str(head_dir / "probe_head.safetensors"), ) (head_dir / "io.yaml").write_text("name: safety\nmodel: ~\n") manifest = root / "adapters.yaml" @@ -190,21 +190,21 @@ def _write_classifier_manifest(root): "safety:\n" f" path: {head_dir}\n" " type: lora\n" - " kind: classifier\n" - f" labels: [{', '.join(_CLASSIFIER_LABELS)}]\n" + " kind: probe\n" + f" labels: [{', '.join(_PROBE_LABELS)}]\n" ) return manifest -def _compose_classifier(root, enable_audio): - manifest = _write_classifier_manifest(root) +def _compose_probe(root, enable_audio): + manifest = _write_probe_manifest(root) output_dir = root / "model" cmd = [ sys.executable, "-m", "granite_switch.composer.compose_granite_switch", "--base-model", - _CLASSIFIER_BASE, + _PROBE_BASE, "--adapters", str(manifest), "--output", @@ -214,7 +214,7 @@ def _compose_classifier(root, enable_audio): cmd.append("--enable-audio") result = subprocess.run(cmd, capture_output=True, text=True, timeout=BUILD_TIMEOUT) assert result.returncode == 0, ( - f"classifier compose failed (exit {result.returncode}).\n" + f"probe compose failed (exit {result.returncode}).\n" f"STDOUT tail:\n{result.stdout[-2000:]}\n" f"STDERR tail:\n{result.stderr[-2000:]}" ) @@ -222,21 +222,19 @@ def _compose_classifier(root, enable_audio): @pytest.fixture(scope="module", params=[True, False], ids=["audio", "text"]) -def classifier_checkpoint(request, tmp_path_factory): - root = tmp_path_factory.mktemp( - "classifier-audio" if request.param else "classifier" - ) - return request.param, _compose_classifier(root, enable_audio=request.param) +def probe_checkpoint(request, tmp_path_factory): + root = tmp_path_factory.mktemp("probe-audio" if request.param else "probe") + return request.param, _compose_probe(root, enable_audio=request.param) @pytest.mark.xdist_group("multi_audio_compose_e2e") -def test_classifier_checkpoint_loads_strictly_and_routes(classifier_checkpoint): +def test_probe_checkpoint_loads_strictly_and_routes(probe_checkpoint): torch = pytest.importorskip("torch") from transformers import AutoTokenizer from granite_switch.hf import GraniteSwitchForCausalLM - enable_audio, output_dir = classifier_checkpoint + enable_audio, output_dir = probe_checkpoint config = json.loads((output_dir / "config.json").read_text()) if enable_audio: # Guard the guard: the marker must push the vocabulary past the controls. @@ -262,9 +260,9 @@ def test_classifier_checkpoint_loads_strictly_and_routes(classifier_checkpoint): assert input_ids[0, -1].item() == marker with torch.no_grad(): out = model(input_ids=input_ids, use_cache=False) - assert model.model._last_classifier_indices[0, -1].item() == 2 + assert model.model._last_probe_indices[0, -1].item() == 2 assert not model.model._last_lora_indices.any() - label_ids = model.config.classifier_label_token_ids[1] + label_ids = model.config.probe_label_token_ids[1] assert int(out.logits[0, -1].argmax()) == label_ids[1] diff --git a/tests/composer/test_classifier_compose.py b/tests/composer/test_probe_compose.py similarity index 79% rename from tests/composer/test_classifier_compose.py rename to tests/composer/test_probe_compose.py index c4136bc9..0af60d29 100644 --- a/tests/composer/test_classifier_compose.py +++ b/tests/composer/test_probe_compose.py @@ -1,31 +1,31 @@ # SPDX-License-Identifier: Apache-2.0 -"""Compose-time loading of trained classifier-head weights. +"""Compose-time loading of trained probe-head weights. -A "classifier" slot is the classifier-head analog of a LoRA adapter: it shares +A "probe" slot is the probe-head analog of a LoRA adapter: it shares the same control-token / adapter-index machinery, but instead of LoRA weights it carries a small linear head (``weight [n_labels, hidden]`` + ``bias [n_labels]``) -that emits a per-request verdict. Each classifier slot may declare a DIFFERENT -number of labels; the head bank is padded to ``max_classifier_labels`` (like the +that emits a per-request verdict. Each probe slot may declare a DIFFERENT +number of labels; the head bank is padded to ``max_probe_labels`` (like the LoRA bank pads ranks to ``max_lora_rank``), so a slot's trained head lands in the top ``n_labels`` rows of its bank row and the pad rows stay zero. The compose -pipeline builds the ``SwitchedClassifierHead`` from config and transfers the -trained head weights from each classifier directory's -``classifier_head.safetensors`` into bank row ``slot = adapter_index - 1``. +pipeline builds the ``SwitchedProbeHead`` from config and transfers the +trained head weights from each probe directory's +``probe_head.safetensors`` into bank row ``slot = adapter_index - 1``. These tests build the model through :class:`GraniteSwitchComposer` (per CLAUDE.md rule #5 — never hand-assemble the config) against the real ``granite-4.0-micro`` -base, with a *synthetic* classifier head whose weights make one label win +base, with a *synthetic* probe head whose weights make one label win deterministically, and assert: -1. ``classifier_head.weight[slot, :n_labels]`` / ``.bias[slot, :n_labels]`` exactly +1. ``probe_head.weight[slot, :n_labels]`` / ``.bias[slot, :n_labels]`` exactly equal the synthetic tensors, and the pad rows above ``n_labels`` stay zero. -2. A classifier control-token forward emits the winning label's token id (the HF +2. A probe control-token forward emits the winning label's token id (the HF verdict exit scatters the head's argmax onto that slot's - ``classifier_label_token_ids``). -3. In a mixed ``[lora, classifier]`` compose the LoRA slot's bank row is populated - while the classifier slot's LoRA row stays zero (compose orders LoRAs first, so - slot numbering is stable), and the classifier row lands in its slot. -4. Two classifier slots with DIFFERENT label counts (2 and 3) each land padded + ``probe_label_token_ids``). +3. In a mixed ``[lora, probe]`` compose the LoRA slot's bank row is populated + while the probe slot's LoRA row stays zero (compose orders LoRAs first, so + slot numbering is stable), and the probe row lands in its slot. +4. Two probe slots with DIFFERENT label counts (2 and 3) each land padded correctly and each emit a token id from their own label set. Marked slow + requires_model: composing the ~3B micro base is expensive, but it is @@ -41,7 +41,7 @@ import granite_switch.hf # noqa: F401 — registers AutoModel/AutoConfig from granite_switch.composer.compose_utils import GraniteSwitchComposer -from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE +from granite_switch.composer.weight_transfer import PROBE_HEAD_FILE BASE_MODEL = "ibm-granite/granite-4.0-micro" # The mixed test embeds a real LoRA, which must match its base. The RAG library @@ -77,7 +77,7 @@ def _hidden_size(model=BASE_MODEL): def _single_token_labels(tokenizer, candidates, count=2): """Pick ``count`` label words that each encode to exactly one token id. - The classifier verdict exit writes at one logit position per request, so a + The probe verdict exit writes at one logit position per request, so a label must be a single token. Returns ``(names, ids)``. """ names, ids = [], [] @@ -92,18 +92,18 @@ def _single_token_labels(tokenizer, candidates, count=2): def _write_synthetic_head(directory, weight, bias, labels): - """Write a classifier slot dir: classifier_head.safetensors (+ config).""" + """Write a probe slot dir: probe_head.safetensors (+ config).""" from safetensors.torch import save_file os.makedirs(directory, exist_ok=True) save_file( {"weight": weight.contiguous(), "bias": bias.contiguous()}, - os.path.join(directory, CLASSIFIER_HEAD_FILE), + os.path.join(directory, PROBE_HEAD_FILE), ) # Optional manifest metadata — the composer reads kinds/labels from kwargs, - # not this file, but a real classifier dir carries it, so include it. + # not this file, but a real probe dir carries it, so include it. with open(os.path.join(directory, "adapter_config.json"), "w") as f: - json.dump({"kind": "classifier", "labels": list(labels)}, f) + json.dump({"kind": "probe", "labels": list(labels)}, f) def _winning_head(num_labels, hidden, winner, dtype=torch.float32): @@ -122,7 +122,7 @@ def _winning_head(num_labels, hidden, winner, dtype=torch.float32): return weight, bias -def _classifier_prompt_ids(tokenizer, control_token_id): +def _probe_prompt_ids(tokenizer, control_token_id): """`` as a batch of input ids for a detect forward. The control token (marker) is the last token, as in a chat-template render, @@ -172,8 +172,8 @@ def _find_rag_lora_dir(): # --------------------------------------------------------------------------- # # Tests # --------------------------------------------------------------------------- # -def test_classifier_only_compose_loads_trained_head(tmp_path): - """A single classifier slot: weights land in the bank and the verdict fires.""" +def test_probe_only_compose_loads_trained_head(tmp_path): + """A single probe slot: weights land in the bank and the verdict fires.""" tokenizer = _tokenizer() hidden = _hidden_size() @@ -187,7 +187,7 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): clf_dir = str(tmp_path / "detect") _write_synthetic_head(clf_dir, weight, bias, label_names) - # A control token id for the single classifier slot. Any spare vocab id + # A control token id for the single probe slot. Any spare vocab id # works for construction; the switch fires on this id, which the prompt # helper places last, as the chat template does. control_id = tokenizer.vocab_size - 5 @@ -196,18 +196,18 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): base_model_name_or_path=BASE_MODEL, adapter_paths=[clf_dir], adapter_token_ids=[control_id], - adapter_substitute_token_ids=[control_id], # inert for a classifier slot + adapter_substitute_token_ids=[control_id], # inert for a probe slot adapter_names=["detect"], - adapter_kinds=["classifier"], - # Full-width per-slot label count (single classifier slot -> [num_labels]); - # ragged label metadata is one list per classifier slot. - classifier_label_token_ids=[label_ids], + adapter_kinds=["probe"], + # Full-width per-slot label count (single probe slot -> [num_labels]); + # ragged label metadata is one list per probe slot. + probe_label_token_ids=[label_ids], ) model.eval() # (1) The trained head is in bank row slot = index - 1 = 0, top n_labels rows. - head = model.model.classifier_head - assert head is not None, "classifier_head was not built" + head = model.model.probe_head + assert head is not None, "probe_head was not built" torch.testing.assert_close( head.weight[0, :num_labels].float(), weight.to(head.weight.dtype).float() ) @@ -216,7 +216,7 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): ) # (2) A detect forward emits the winning label's token id. - input_ids = _classifier_prompt_ids(tokenizer, control_id).to( + input_ids = _probe_prompt_ids(tokenizer, control_id).to( next(model.parameters()).device ) with torch.no_grad(): @@ -228,8 +228,8 @@ def test_classifier_only_compose_loads_trained_head(tmp_path): ) -def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): - """[lora, classifier]: LoRA slot 0 populated, classifier slot 1 loaded.""" +def test_mixed_lora_and_probe_preserves_slot_numbering(tmp_path): + """[lora, probe]: LoRA slot 0 populated, probe slot 1 loaded.""" tokenizer = _tokenizer(MIXED_BASE) hidden = _hidden_size(MIXED_BASE) lora_dir = _find_rag_lora_dir() @@ -253,54 +253,50 @@ def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): adapter_token_ids=[lora_ctrl, clf_ctrl], adapter_substitute_token_ids=[lora_ctrl, clf_ctrl], adapter_names=["rag_lora", "detect"], - adapter_kinds=["lora", "classifier"], - # Full width: None on the LoRA slot, ids on the classifier slot. - classifier_label_token_ids=[None, label_ids], + adapter_kinds=["lora", "probe"], + # Full width: None on the LoRA slot, ids on the probe slot. + probe_label_token_ids=[None, label_ids], ) model.eval() - head = model.model.classifier_head + head = model.model.probe_head assert head is not None - # Classifier slot is index 1 -> bank row 1; slot 0 (the LoRA slot) is never - # touched by the classifier transfer, so it stays at the bank's zero init. + # Probe slot is index 1 -> bank row 1; slot 0 (the LoRA slot) is never + # touched by the probe transfer, so it stays at the bank's zero init. torch.testing.assert_close( head.weight[1, :num_labels].float(), weight.to(head.weight.dtype).float() ) torch.testing.assert_close( head.bias[1, :num_labels].float(), bias.to(head.bias.dtype).float() ) - # Slot 0's classifier row must stay all-zero (never loaded) — a stronger check + # Slot 0's probe row must stay all-zero (never loaded) — a stronger check # than "not equal to the synthetic head" now that the bank is zero-init. - assert head.weight[0].abs().sum() == 0, ( - "the LoRA slot's classifier row was overwritten" - ) - assert head.bias[0].abs().sum() == 0, ( - "the LoRA slot's classifier bias was overwritten" - ) + assert head.weight[0].abs().sum() == 0, "the LoRA slot's probe row was overwritten" + assert head.bias[0].abs().sum() == 0, "the LoRA slot's probe bias was overwritten" - # The LoRA bank row for slot 0 must be populated, while the classifier + # The LoRA bank row for slot 0 must be populated, while the probe # slot's LoRA row (slot 1) must stay unloaded. lora_B is the discriminator: # it is ZERO-initialized and only becomes nonzero when an adapter is loaded # (lora_A is kaiming-initialized, so an unloaded row is nonzero there too and # can't distinguish loaded from unloaded). So: slot 0 lora_B nonzero (loaded), - # slot 1 lora_B all-zero (classifier slot, never received LoRA weights). + # slot 1 lora_B all-zero (probe slot, never received LoRA weights). lora_B = _first_lora_b_tensor(model) assert lora_B is not None, "no stacked LoRA B tensor found on the model" assert lora_B[0].abs().sum() > 0, "LoRA slot 0 bank row is unexpectedly all-zero" assert lora_B[1].abs().sum() == 0, ( - "classifier slot 1 LoRA row is not zero — slot numbering leaked" + "probe slot 1 LoRA row is not zero — slot numbering leaked" ) - # A detect forward on the classifier control token emits the winning label. + # A detect forward on the probe control token emits the winning label. device = next(model.parameters()).device - input_ids = _classifier_prompt_ids(tokenizer, clf_ctrl).to(device) + input_ids = _probe_prompt_ids(tokenizer, clf_ctrl).to(device) with torch.no_grad(): out = model(input_ids=input_ids, use_cache=False) emitted = int(out.logits[0, -1].argmax().item()) assert emitted == label_ids[winner] assert model.model._last_lora_indices[0, -1].item() == 0 - assert model.model._last_classifier_indices[0, -1].item() == 2 + assert model.model._last_probe_indices[0, -1].item() == 2 # The LoRA slot is selected on its own and leaves the LM logits intact. lora_ids = torch.cat( @@ -309,12 +305,12 @@ def test_mixed_lora_and_classifier_preserves_slot_numbering(tmp_path): with torch.no_grad(): out = model(input_ids=lora_ids, use_cache=False) assert model.model._last_lora_indices[0, -1].item() == 1 - assert model.model._last_classifier_verdict is None + assert model.model._last_probe_verdict is None assert torch.isfinite(out.logits[0, -1]).all() def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): - """Two classifier slots with DIFFERENT label counts (2 and 3). + """Two probe slots with DIFFERENT label counts (2 and 3). Each head lands in the top rows of its (padded) bank row, the pad rows above its own label count stay zero, and a detect forward on each slot's control @@ -354,12 +350,12 @@ def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): adapter_token_ids=[ctrl_a, ctrl_b], adapter_substitute_token_ids=[ctrl_a, ctrl_b], adapter_names=["detect_a", "detect_b"], - adapter_kinds=["classifier", "classifier"], - classifier_label_token_ids=[ids_a, ids_b], + adapter_kinds=["probe", "probe"], + probe_label_token_ids=[ids_a, ids_b], ) model.eval() - head = model.model.classifier_head + head = model.model.probe_head assert head is not None # Bank is padded to max_labels; each slot's head fills its top n rows. assert head.weight.shape[1] == max_labels @@ -385,7 +381,7 @@ def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): # Slot A fires -> emits one of A's ids (never a B id, never a padded slot). with torch.no_grad(): out_a = model( - input_ids=_classifier_prompt_ids(tokenizer, ctrl_a).to(device), + input_ids=_probe_prompt_ids(tokenizer, ctrl_a).to(device), use_cache=False, ) emitted_a = int(out_a.logits[0, -1].argmax().item()) @@ -395,7 +391,7 @@ def test_ragged_label_counts_pad_and_fire_per_slot(tmp_path): # Slot B fires -> emits one of B's ids, including its 3rd (only-in-B) label. with torch.no_grad(): out_b = model( - input_ids=_classifier_prompt_ids(tokenizer, ctrl_b).to(device), + input_ids=_probe_prompt_ids(tokenizer, ctrl_b).to(device), use_cache=False, ) emitted_b = int(out_b.logits[0, -1].argmax().item()) @@ -406,7 +402,7 @@ def _first_lora_b_tensor(model): """Return the first stacked LoRA-B bank tensor [num_adapters, 1, out, r]. lora_B is zero-init and only nonzero once an adapter is loaded, so it cleanly - distinguishes a loaded LoRA slot from an unloaded (classifier) slot — unlike + distinguishes a loaded LoRA slot from an unloaded (probe) slot — unlike lora_A, which is kaiming-init and nonzero even when never loaded. """ for name, buf in model.named_parameters(): @@ -426,8 +422,8 @@ def _first_lora_b_tensor(model): "ibm-granite/granite-4.2-3b", # ChatML family ], ) -def test_classifier_drops_generation_prompt(base): - """A classifier render ignores ``add_generation_prompt=True``. +def test_probe_drops_generation_prompt(base): + """A probe render ignores ``add_generation_prompt=True``. The verdict is read at the marker, so the assistant generation prompt must never follow it. vLLM's chat endpoint defaults the flag to True, so the @@ -439,7 +435,7 @@ def test_classifier_drops_generation_prompt(base): ) tokenizer = _tokenizer(base) - discovered = [(None, "safety", "classifier", None)] + discovered = [(None, "safety", "probe", None)] add_control_tokens(tokenizer, discovered) configure_chat_template(tokenizer, discovered) messages = [{"role": "user", "content": "Is this text safe?"}] @@ -462,7 +458,7 @@ def render(gen, **kwargs): "ibm-granite/granite-4.2-3b", # ChatML family ], ) -def test_classifier_marker_is_last_token(base): +def test_probe_marker_is_last_token(base): """The marker is the very last token, after the last turn's close. Everything before it is byte-identical to a no-adapter render without a @@ -474,7 +470,7 @@ def test_classifier_marker_is_last_token(base): ) tokenizer = _tokenizer(base) - discovered = [(None, "safety", "classifier", None)] + discovered = [(None, "safety", "probe", None)] (control_id,), _ = add_control_tokens(tokenizer, discovered) configure_chat_template(tokenizer, discovered) messages = [ @@ -505,13 +501,13 @@ def test_classifier_marker_is_last_token(base): ("ibm-granite/granite-4.2-3b", "<|im_start|>"), ], ) -def test_classifier_substitute_is_the_bases_role_open_marker(base, role_open): +def test_probe_substitute_is_the_bases_role_open_marker(base, role_open): from granite_switch.composer.compose_granite_switch import ( build_control_token_lists, ) tokenizer = _tokenizer(base) - discovered = [(None, "safety", "classifier", None)] + discovered = [(None, "safety", "probe", None)] _, _, substitutes = build_control_token_lists( tokenizer, discovered, base_reset=False ) diff --git a/tests/composer/test_sr_compose_e2e.py b/tests/composer/test_sr_compose_e2e.py index d9f9c316..a32c7bbf 100644 --- a/tests/composer/test_sr_compose_e2e.py +++ b/tests/composer/test_sr_compose_e2e.py @@ -184,55 +184,55 @@ def _create_sr_adapter( save_file(state_dict, str(path / "adapter_model.safetensors")) -def _compose_sr_classifier(path: Path, read_layers=(None,)): - """Tiny SR + classifier checkpoint, also used by the vLLM test. +def _compose_sr_probe(path: Path, read_layers=(None,)): + """Tiny SR + probe checkpoint, also used by the vLLM test. An identity head exposes every component of its input, so its scores pin the exact layer and normalization read without random margins. """ from granite_switch.composer import GraniteSwitchComposer - from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE + from granite_switch.composer.weight_transfer import PROBE_HEAD_FILE - base, sr, classifier = (path / name for name in ("base", "sr", "classifier")) + base, sr, probe = (path / name for name in ("base", "sr", "probe")) base.mkdir(parents=True) - classifier.mkdir() + probe.mkdir() _create_base_model(base) _create_sr_adapter(sr) save_file( {"weight": torch.eye(HIDDEN), "bias": torch.zeros(HIDDEN)}, - str(classifier / CLASSIFIER_HEAD_FILE), + str(probe / PROBE_HEAD_FILE), ) return GraniteSwitchComposer.from_base_and_adapters( base_model_name_or_path=str(base), - adapter_paths=[str(sr), *[str(classifier)] * len(read_layers)], + adapter_paths=[str(sr), *[str(probe)] * len(read_layers)], adapter_token_ids=list( range(CONTROL_TOKEN_ID, CONTROL_TOKEN_ID + 1 + len(read_layers)) ), adapter_substitute_token_ids=[1, *[2] * len(read_layers)], adapter_names=["sr", *[f"detect{i}" for i in range(len(read_layers))]], - adapter_kinds=["lora", *["classifier"] * len(read_layers)], - classifier_label_token_ids=[ + adapter_kinds=["lora", *["probe"] * len(read_layers)], + probe_label_token_ids=[ None, *[list(range(100, 100 + HIDDEN))] * len(read_layers), ], - classifier_read_layers=[None, *read_layers], + probe_read_layers=[None, *read_layers], switch_head_dim=16, ).eval() @pytest.mark.parametrize("read_layer", [None, 0, NUM_LAYERS - 1]) -def test_sr_and_classifier_slots_select_separately(tmp_path, capsys, read_layer): +def test_sr_and_probe_slots_select_separately(tmp_path, capsys, read_layer): """Identity scores prove layer and normalization before/after reload. - The classifier marker routes out of the SR stream, so the probe reads the + The probe marker routes out of the SR stream, so the probe reads the base stream at its layer; an SR request in the same checkpoint still adapts. """ from granite_switch.hf import load_model - model = _compose_sr_classifier(tmp_path, read_layers=[read_layer]) + model = _compose_sr_probe(tmp_path, read_layers=[read_layer]) assert "cross_stream slots are all-zero" not in capsys.readouterr().out assert model.config.dual_stream - head = model.model.classifier_head + head = model.model.probe_head torch.testing.assert_close(head.weight[1], torch.eye(HIDDEN)) assert torch.count_nonzero(head.weight[0]) == 0 for layer in model.model.layers: @@ -254,21 +254,21 @@ def run(m, ids): ) return output, base, adapted - labels = model.config.classifier_label_token_ids[1] + labels = model.config.probe_label_token_ids[1] non_labels = torch.ones(VOCAB_SIZE, dtype=torch.bool) non_labels[labels] = False probe_ids = torch.tensor([[10, 21, 20, CONTROL_TOKEN_ID + 1]]) probe, base, _ = run(model, probe_ids) assert model.model._last_lora_indices.tolist() == [[0, 0, 0, 0]] - assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] + assert model.model._last_probe_indices.tolist() == [[0, 0, 0, 2]] torch.testing.assert_close(probe.logits[0, -1, labels], base) assert torch.isneginf(probe.logits[0, -1, non_labels]).all() sr_ids = torch.tensor([[10, CONTROL_TOKEN_ID, 20, 30]]) sr, base, adapted = run(model, sr_ids) assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1]] - assert not model.model._last_classifier_indices.any() + assert not model.model._last_probe_indices.any() assert not torch.allclose(adapted, base), "fixture must exercise SR adaptation" assert torch.isfinite(sr.logits[0, -1]).all() diff --git a/tests/composer/test_tokenizer_setup.py b/tests/composer/test_tokenizer_setup.py index 2172ff80..c2c0ca23 100644 --- a/tests/composer/test_tokenizer_setup.py +++ b/tests/composer/test_tokenizer_setup.py @@ -344,19 +344,19 @@ def test_builtin_technology_uses_the_probed_substitute(self): assert subs == [42] - def test_classifier_uses_its_own_substitute(self): - adapters = [("/b", "code", "lora", None), (None, "safety", "classifier", None)] + def test_probe_uses_its_own_substitute(self): + adapters = [("/b", "code", "lora", None), (None, "safety", "probe", None)] subs = build_substitute_token_ids( - adapters, lora_substitute_id=42, classifier_substitute_id=9 + adapters, lora_substitute_id=42, probe_slot_substitute_id=9 ) assert subs == [42, 9] - def test_classifier_without_its_substitute_raises(self): - adapters = [(None, "safety", "classifier", None)] + def test_probe_without_its_substitute_raises(self): + adapters = [(None, "safety", "probe", None)] - with pytest.raises(ValueError, match="classifier_substitute_id"): + with pytest.raises(ValueError, match="probe_slot_substitute_id"): build_substitute_token_ids(adapters, lora_substitute_id=42) @@ -641,7 +641,7 @@ def convert_ids_to_tokens(self, ids): class TestResolveLabelTokenIds: - """Tests for resolve_label_token_ids (classifier label word -> token id).""" + """Tests for resolve_label_token_ids (probe label word -> token id).""" def test_single_token_labels_resolve_in_order(self): tok = _LabelTokenizer({"safe": 19193, "unsafe": 39257}) diff --git a/tests/hf/test_classifier.py b/tests/hf/test_probe.py similarity index 65% rename from tests/hf/test_classifier.py rename to tests/hf/test_probe.py index e0fa9df1..51abf06a 100644 --- a/tests/hf/test_classifier.py +++ b/tests/hf/test_probe.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -"""Classifier-substitute path tests for GraniteSwitchForCausalLM (HF backend). +"""Probe-substitute path tests for GraniteSwitchForCausalLM (HF backend). -A classifier slot is an alternative to a LoRA adapter: it shares the adapter -index/control-token machinery, the switch splits classifier positions out of -the LoRA stream, and the classifier head reads its slot's read layer (the final +A probe slot is an alternative to a LoRA adapter: it shares the adapter +index/control-token machinery, the switch splits probe positions out of +the LoRA stream, and the probe head reads its slot's read layer (the final hidden state by default) to emit ONE verdict per request. The verdict is read at the control token, which must be the request's last token. The verdict exits as a GENERATED LABEL WORD, matching the vLLM backend: the LM head rewrites the marker's logit row so the vocab is -inf except at -``classifier_label_token_ids``, where the per-label verdict scores are placed. There is no separate float ``classifier_logits`` output field. +``probe_label_token_ids``, where the per-label verdict scores are placed. There is no separate float ``probe_logits`` output field. CPU, random weights, no checkpoint — mirrors tests/hf/test_model_forward.py. """ @@ -25,13 +25,13 @@ def _set_adapter_token_ids(model, token_ids): model.model.adapter_token_ids.data = torch.tensor(token_ids, dtype=torch.long) -def _classifier_config( +def _probe_config( adapter_kinds, num_labels=2, label_token_ids=(100, 200), ): # Full-width label token ids (length num_adapters, None on LoRA slots). Every - # classifier slot gets the same ids here; mixed counts are covered by the + # probe slot gets the same ids here; mixed counts are covered by the # compose tests. label_token_ids = list(label_token_ids) return GraniteSwitchConfig( @@ -46,8 +46,8 @@ def _classifier_config( adapter_substitute_token_ids=[1, 1], adapter_names=["adapter_1", "adapter_2"], adapter_kinds=adapter_kinds, - classifier_label_token_ids=[ - list(label_token_ids) if k == "classifier" else None for k in adapter_kinds + probe_label_token_ids=[ + list(label_token_ids) if k == "probe" else None for k in adapter_kinds ], max_lora_rank=4, adapter_ranks=[4, 4], @@ -55,21 +55,21 @@ def _classifier_config( ) -def test_classifier_head_built_only_when_slot_present(): - m_lora = GraniteSwitchForCausalLM(_classifier_config(["lora", "lora"])) - m_clf = GraniteSwitchForCausalLM(_classifier_config(["lora", "classifier"])) - assert m_lora.model.classifier_head is None - assert m_clf.model.classifier_head is not None +def test_probe_head_built_only_when_slot_present(): + m_lora = GraniteSwitchForCausalLM(_probe_config(["lora", "lora"])) + m_clf = GraniteSwitchForCausalLM(_probe_config(["lora", "probe"])) + assert m_lora.model.probe_head is None + assert m_clf.model.probe_head is not None # Bank sized to the full index space so slot = index - 1 addresses it. - assert m_clf.model.classifier_head.num_classifier_slots == 2 - assert m_clf.model.classifier_head.num_labels == 2 + assert m_clf.model.probe_head.num_probe_slots == 2 + assert m_clf.model.probe_head.num_labels == 2 -def test_classifier_rewrites_last_logit_row_to_label_word(): - # adapter 2 (token 251) is a classifier. Its verdict rides the LM logits: +def test_probe_rewrites_last_logit_row_to_label_word(): + # adapter 2 (token 251) is a probe. Its verdict rides the LM logits: # the last prompt row is -inf everywhere except at the label token ids, # where the per-label verdict scores land — so decoding emits a label word. - config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + config = _probe_config(["lora", "probe"], label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() @@ -80,8 +80,8 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): out = model(input_ids=input_ids) last_row = out.logits[0, -1, :] # [vocab] - # Single classifier slot -> classifier position 0; its own label ids. - label_ids = config.classifier_label_token_ids[1] # slot 2 -> index 1 + # Single probe slot -> probe position 0; its own label ids. + label_ids = config.probe_label_token_ids[1] # slot 2 -> index 1 n = len(label_ids) non_label = [t for t in range(config.vocab_size) if t not in label_ids] @@ -89,16 +89,16 @@ def test_classifier_rewrites_last_logit_row_to_label_word(): assert torch.isinf(last_row[non_label]).all() assert (last_row[non_label] < 0).all() assert torch.isfinite(last_row[torch.tensor(label_ids)]).all() - # The label-id scores equal the classifier head's per-label verdict, sliced + # The label-id scores equal the probe head's per-label verdict, sliced # to this slot's real label count (padded columns are never scattered). - verdict = model.model._last_classifier_verdict[2][0, :n] # [n_labels] + verdict = model.model._last_probe_verdict[2][0, :n] # [n_labels] assert torch.allclose(last_row[torch.tensor(label_ids)], verdict.to(last_row.dtype)) # The argmax (emitted token) is one of the label ids. assert int(last_row.argmax()) in label_ids -def test_lora_and_classifier_slots_select_separately(): - config = _classifier_config(["lora", "classifier"]) +def test_lora_and_probe_slots_select_separately(): + config = _probe_config(["lora", "probe"]) model = GraniteSwitchForCausalLM(config).eval() _set_adapter_token_ids(model, config.adapter_token_ids) @@ -106,17 +106,17 @@ def test_lora_and_classifier_slots_select_separately(): with torch.no_grad(): model(input_ids=torch.tensor([[10, 250, 20, 30]])) assert model.model._last_lora_indices.tolist() == [[0, 1, 1, 1]] - assert not model.model._last_classifier_indices.any() - assert model.model._last_classifier_verdict is None + assert not model.model._last_probe_indices.any() + assert model.model._last_probe_verdict is None - # The classifier marker selects slot 2 and routes out of the LoRA stream. + # The probe marker selects slot 2 and routes out of the LoRA stream. input_ids = torch.tensor([[10, 20, 30, 251]]) with torch.no_grad(): output = model(input_ids=input_ids) assert model.model._last_lora_indices.tolist() == [[0, 0, 0, 0]] - assert model.model._last_classifier_indices.tolist() == [[0, 0, 0, 2]] - label_ids = config.classifier_label_token_ids[1] + assert model.model._last_probe_indices.tolist() == [[0, 0, 0, 2]] + label_ids = config.probe_label_token_ids[1] last_logits = output.logits[0, -1] assert torch.isfinite(last_logits[torch.tensor(label_ids)]).all() assert torch.isneginf( @@ -124,11 +124,11 @@ def test_lora_and_classifier_slots_select_separately(): ).all() -def test_non_classifier_request_logits_untouched(): - # No classifier token (251) present -> last token is base (index 0), the +def test_non_probe_request_logits_untouched(): + # No probe token (251) present -> last token is base (index 0), the # head's last-token verdict row is zero, and NO row is blanked: a - # non-classifier request keeps a full finite vocab row. - config = _classifier_config(["lora", "classifier"]) + # non-probe request keeps a full finite vocab row. + config = _probe_config(["lora", "probe"]) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() @@ -137,33 +137,33 @@ def test_non_classifier_request_logits_untouched(): with torch.no_grad(): out = model(input_ids=input_ids) - assert (model.model._last_classifier_indices == 0).all() + assert (model.model._last_probe_indices == 0).all() # Full finite vocab row — nothing blanked. assert torch.isfinite(out.logits[0, -1, :]).all() -def test_mixed_batch_only_classifier_row_rewritten(): - # Batch: row 0 is a classifier request (251 fires), row 1 is plain LM. - config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) +def test_mixed_batch_only_probe_row_rewritten(): + # Batch: row 0 is a probe request (251 fires), row 1 is plain LM. + config = _probe_config(["lora", "probe"], label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() input_ids = torch.tensor( [ - [10, 20, 30, 40, 251], # classifier + [10, 20, 30, 40, 251], # probe [10, 20, 30, 40, 50], # plain ] ) with torch.no_grad(): out = model(input_ids=input_ids) - label_ids = torch.tensor(config.classifier_label_token_ids[1]) # slot 2 + label_ids = torch.tensor(config.probe_label_token_ids[1]) # slot 2 non_label = [ t for t in range(config.vocab_size) if t not in set(label_ids.tolist()) ] - # Classifier row rewritten; plain row fully finite. + # Probe row rewritten; plain row fully finite. assert torch.isinf(out.logits[0, -1, non_label]).all() assert torch.isfinite(out.logits[1, -1, :]).all() @@ -171,7 +171,7 @@ def test_mixed_batch_only_classifier_row_rewritten(): def test_decode_step_after_verdict_is_plain_lm(): # After the verdict token, a cached decode step carries no marker, so it runs # as a plain LM step: no verdict, full finite vocab row. - config = _classifier_config(["lora", "classifier"]) + config = _probe_config(["lora", "probe"]) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() @@ -179,7 +179,7 @@ def test_decode_step_after_verdict_is_plain_lm(): input_ids = torch.tensor([[10, 20, 30, 251]]) with torch.no_grad(): out = model(input_ids=input_ids, use_cache=True) - assert model.model._last_classifier_verdict is not None + assert model.model._last_probe_verdict is not None with torch.no_grad(): step = model( @@ -187,15 +187,15 @@ def test_decode_step_after_verdict_is_plain_lm(): past_key_values=out.past_key_values, use_cache=True, ) - assert model.model._last_classifier_verdict is None + assert model.model._last_probe_verdict is None assert torch.isfinite(step.logits).all() -def test_classifier_missing_label_token_ids_fails_loud(): - # A classifier slot with no resolved label token ids cannot emit a label word. +def test_probe_missing_label_token_ids_fails_loud(): + # A probe slot with no resolved label token ids cannot emit a label word. # Rejected when the config is built, so a bad checkpoint fails at load rather # than at the first inference. - with pytest.raises(ValueError, match="classifier_label_token_ids"): + with pytest.raises(ValueError, match="probe_label_token_ids"): GraniteSwitchConfig( vocab_size=300, hidden_size=64, @@ -207,8 +207,8 @@ def test_classifier_missing_label_token_ids_fails_loud(): adapter_token_ids=[250, 251], adapter_substitute_token_ids=[1, 1], adapter_names=["adapter_1", "adapter_2"], - adapter_kinds=["lora", "classifier"], - classifier_label_token_ids=None, + adapter_kinds=["lora", "probe"], + probe_label_token_ids=None, max_lora_rank=4, adapter_ranks=[4, 4], switch_head_dim=16, @@ -216,9 +216,9 @@ def test_classifier_missing_label_token_ids_fails_loud(): def test_lora_only_path_unaffected(): - # A model with no classifier slot behaves exactly as before: no classifier + # A model with no probe slot behaves exactly as before: no probe # head, no verdict, full finite LM logits. - config = _classifier_config(["lora", "lora"]) + config = _probe_config(["lora", "lora"]) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() @@ -227,28 +227,28 @@ def test_lora_only_path_unaffected(): with torch.no_grad(): out = model(input_ids=input_ids) - assert model.model.classifier_head is None - assert model.model._last_classifier_verdict is None + assert model.model.probe_head is None + assert model.model._last_probe_verdict is None assert torch.isfinite(out.logits).all() assert out.logits.shape == (1, 6, config.vocab_size) def test_lora_generation_on_mixed_model_decodes(): - # A model composed WITH a classifier slot also serves plain LoRA/base - # generation. Firing the LoRA (adapter 1, token 250) — not the classifier — + # A model composed WITH a probe slot also serves plain LoRA/base + # generation. Firing the LoRA (adapter 1, token 250) — not the probe — # must decode multiple tokens with no verdict and full LM logits. - config = _classifier_config(["lora", "classifier"]) + config = _probe_config(["lora", "probe"]) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() - assert model.model.classifier_head is not None # mixed model has the bank + assert model.model.probe_head is not None # mixed model has the bank - # Prefill fires the LoRA (250), never the classifier (251). + # Prefill fires the LoRA (250), never the probe (251). input_ids = torch.tensor([[10, 250, 20, 30, 40]]) with torch.no_grad(): out = model(input_ids=input_ids, use_cache=True) - # No classifier fired -> no verdict, full finite vocab row. - assert model.model._last_classifier_verdict is None + # No probe fired -> no verdict, full finite vocab row. + assert model.model._last_probe_verdict is None assert torch.isfinite(out.logits).all() # Cached decode step of the same (LoRA) request yields finite logits. @@ -260,25 +260,25 @@ def test_lora_generation_on_mixed_model_decodes(): ) assert step.logits.shape == (1, 1, config.vocab_size) assert torch.isfinite(step.logits).all() - assert model.model._last_classifier_verdict is None + assert model.model._last_probe_verdict is None @pytest.mark.parametrize( "kinds, input_ids", [ # Content after the marker. - (["lora", "classifier"], [[10, 20, 251, 30, 40]]), - # Two different classifier slots, the second one last. - (["classifier", "classifier"], [[250, 10, 20, 30, 251]]), + (["lora", "probe"], [[10, 20, 251, 30, 40]]), + # Two different probe slots, the second one last. + (["probe", "probe"], [[250, 10, 20, 30, 251]]), # The same slot's marker twice. - (["lora", "classifier"], [[251, 10, 20, 30, 251]]), + (["lora", "probe"], [[251, 10, 20, 30, 251]]), ], ids=["content-after", "two-slots", "repeated"], ) def test_marker_that_is_not_the_last_token_raises(kinds, input_ids): """The verdict is read at the last token, so any other marker would be ignored silently; it raises instead.""" - config = _classifier_config(kinds, label_token_ids=(100, 200)) + config = _probe_config(kinds, label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() @@ -291,7 +291,7 @@ def test_marker_that_is_not_the_last_token_raises(kinds, input_ids): def test_right_padded_row_rewrites_the_marker_column(): """With right padding the marker is the row's last real token, not its last column; the verdict lands on the marker's own logit row.""" - config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) + config = _probe_config(["lora", "probe"], label_token_ids=(100, 200)) model = GraniteSwitchForCausalLM(config) _set_adapter_token_ids(model, config.adapter_token_ids) model.eval() @@ -301,7 +301,7 @@ def test_right_padded_row_rewrites_the_marker_column(): with torch.no_grad(): out = model(input_ids=input_ids, attention_mask=attention_mask) - label_ids = config.classifier_label_token_ids[1] + label_ids = config.probe_label_token_ids[1] non_label = [t for t in range(config.vocab_size) if t not in label_ids] assert torch.isinf(out.logits[0, 2, non_label]).all() assert torch.isfinite(out.logits[0, -1, :]).all() # padding untouched @@ -313,22 +313,22 @@ def test_right_padded_row_rewrites_the_marker_column(): model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1) -def _randomize_classifier_head(model): +def _randomize_probe_head(model): """The head bank starts at zero; give it weights so the read layer matters.""" generator = torch.Generator().manual_seed(0) - head = model.model.classifier_head + head = model.model.probe_head with torch.no_grad(): head.weight.copy_(torch.randn(head.weight.shape, generator=generator)) head.bias.copy_(torch.randn(head.bias.shape, generator=generator)) -def test_intermediate_layer_classifier_scores_that_layers_output(): +def test_intermediate_layer_probe_scores_that_layers_output(): """An intermediate read scores the marker row at the configured layer.""" - config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) - config.classifier_read_layers = [None, 0] # slot 2: an intermediate layer + config = _probe_config(["lora", "probe"], label_token_ids=(100, 200)) + config.probe_read_layers = [None, 0] # slot 2: an intermediate layer model = GraniteSwitchForCausalLM(config).eval() _set_adapter_token_ids(model, config.adapter_token_ids) - _randomize_classifier_head(model) + _randomize_probe_head(model) with torch.no_grad(): out = model( input_ids=torch.tensor([[10, 20, 30, 40, 50, 60, 70, 251]]), @@ -336,10 +336,10 @@ def test_intermediate_layer_classifier_scores_that_layers_output(): ) assert model.model._last_lora_indices[0, -1].item() == 0 - assert model.model._last_classifier_indices[0, -1].item() == 2 - rows, pos, verdict, slots, _ = model.model._last_classifier_verdict - expected = model.model.classifier_head(out.hidden_states[1][rows, pos], slots) - final = model.model.classifier_head(out.hidden_states[-1][rows, pos], slots) + assert model.model._last_probe_indices[0, -1].item() == 2 + rows, pos, verdict, slots, _ = model.model._last_probe_verdict + expected = model.model.probe_head(out.hidden_states[1][rows, pos], slots) + final = model.model.probe_head(out.hidden_states[-1][rows, pos], slots) torch.testing.assert_close(verdict, expected) assert not torch.allclose(verdict, final) torch.testing.assert_close(out.logits[0, -1, [100, 200]], verdict[0]) @@ -347,11 +347,11 @@ def test_intermediate_layer_classifier_scores_that_layers_output(): def test_final_decoder_layer_read_uses_pre_norm_output(): """The last decoder layer's residual read precedes the final model norm.""" - config = _classifier_config(["lora", "classifier"], label_token_ids=(100, 200)) - config.classifier_read_layers = [None, 1] + config = _probe_config(["lora", "probe"], label_token_ids=(100, 200)) + config.probe_read_layers = [None, 1] model = GraniteSwitchForCausalLM(config).eval() _set_adapter_token_ids(model, config.adapter_token_ids) - _randomize_classifier_head(model) + _randomize_probe_head(model) captured = [] hook = model.model.layers[1].register_forward_hook( lambda _module, _args, output: captured.append(output[0].detach()) @@ -365,31 +365,31 @@ def test_final_decoder_layer_read_uses_pre_norm_output(): finally: hook.remove() - rows, pos, verdict, slots, _ = model.model._last_classifier_verdict - expected = model.model.classifier_head(captured[0][rows, pos], slots) - final_norm = model.model.classifier_head(out.hidden_states[-1][rows, pos], slots) + rows, pos, verdict, slots, _ = model.model._last_probe_verdict + expected = model.model.probe_head(captured[0][rows, pos], slots) + final_norm = model.model.probe_head(out.hidden_states[-1][rows, pos], slots) torch.testing.assert_close(verdict, expected) assert not torch.allclose(verdict, final_norm) def test_batch_slots_score_at_their_own_read_layers(): - config = _classifier_config(["classifier", "classifier"]) - config.classifier_read_layers = [None, 0] + config = _probe_config(["probe", "probe"]) + config.probe_read_layers = [None, 0] model = GraniteSwitchForCausalLM(config).eval() _set_adapter_token_ids(model, config.adapter_token_ids) - _randomize_classifier_head(model) + _randomize_probe_head(model) with torch.no_grad(): out = model( input_ids=torch.tensor([[10, 20, 250], [10, 20, 251]]), output_hidden_states=True, ) - rows, pos, verdict, slots, _ = model.model._last_classifier_verdict + rows, pos, verdict, slots, _ = model.model._last_probe_verdict assert slots.tolist() == [1, 2] - final_score = model.model.classifier_head( + final_score = model.model.probe_head( out.hidden_states[-1][rows[:1], pos[:1]], slots[:1] ) - intermediate_score = model.model.classifier_head( + intermediate_score = model.model.probe_head( out.hidden_states[1][rows[1:], pos[1:]], slots[1:] ) torch.testing.assert_close(verdict[:1], final_score) diff --git a/tests/shared/classifier_recorder.py b/tests/shared/probe_recorder.py similarity index 81% rename from tests/shared/classifier_recorder.py rename to tests/shared/probe_recorder.py index 290ea297..59c87719 100644 --- a/tests/shared/classifier_recorder.py +++ b/tests/shared/probe_recorder.py @@ -1,28 +1,28 @@ # SPDX-License-Identifier: Apache-2.0 -"""Importable model-RPC callbacks for recording sampled classifier rows. +"""Importable model-RPC callbacks for recording sampled probe rows. PP's multiprocessing executor uses standard pickle, so callbacks must live at module scope. Each worker owns its records; the driver reads them through RPC. """ -def install_classifier_recorder(model): +def install_probe_recorder(model): """Record after compute_logits, outside the compiled/captured forward.""" from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM target = next(m for m in model.modules() if isinstance(m, GraniteSwitchForCausalLM)) records = [] - target._classifier_test_records = records + target._probe_test_records = records original = target.compute_logits config = target.config - head = target.model.classifier_head + head = target.model.probe_head def compute_logits(hidden_states): logits = original(hidden_states) slots = hidden_states[:, -1].round().long() for row in (slots > 0).nonzero(as_tuple=True)[0].tolist(): slot = int(slots[row]) - labels = config.classifier_label_token_ids[slot - 1] + labels = config.probe_label_token_ids[slot - 1] n = len(labels) h = config.hidden_size record = { @@ -31,7 +31,7 @@ def compute_logits(hidden_states): "rewritten_label_logits": logits[row, labels].float().cpu(), } # Only a final-layer head reads the packed hidden-state prefix. - if config.classifier_read_layers[slot - 1] is None: + if config.probe_read_layers[slot - 1] is None: record["served_hidden"] = hidden_states[row, :h].float().cpu() record["loaded_weight"] = head.weight[slot - 1, :n].float().cpu() record["loaded_bias"] = head.bias[slot - 1, :n].float().cpu() @@ -41,12 +41,12 @@ def compute_logits(hidden_states): target.compute_logits = compute_logits -def read_classifier_records(model): +def read_probe_records(model): """Return and clear this worker's rows; non-sampling PP ranks return [].""" from granite_switch.vllm.granite_switch_model import GraniteSwitchForCausalLM target = next(m for m in model.modules() if isinstance(m, GraniteSwitchForCausalLM)) - records = target._classifier_test_records + records = target._probe_test_records result = [ {k: v if k == "slot" else v.tolist() for k, v in r.items()} for r in records ] diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 95d51862..8637fa21 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -261,68 +261,62 @@ def test_lora_targets_include_shared_mlp_groups_when_present(self): # ════════════════════════════════════════════════════════════════════ -# N. classifier_read_layers +# N. probe_read_layers # ════════════════════════════════════════════════════════════════════ -def _classifier_kwargs(num_adapters=2, num_decoder_layers=4, **overrides): - """A LoRA slot(s) then one classifier slot, over ``num_decoder_layers``.""" +def _probe_kwargs(num_adapters=2, num_decoder_layers=4, **overrides): + """A LoRA slot(s) then one probe slot, over ``num_decoder_layers``.""" from granite_switch.config import SWITCH_CACHE_LAYERS base = _valid_kwargs( num_adapters=num_adapters, num_hidden_layers=num_decoder_layers + SWITCH_CACHE_LAYERS, - adapter_kinds=["lora"] * (num_adapters - 1) + ["classifier"], - classifier_label_token_ids=[None] * (num_adapters - 1) + [[100, 200]], + adapter_kinds=["lora"] * (num_adapters - 1) + ["probe"], + probe_label_token_ids=[None] * (num_adapters - 1) + [[100, 200]], ) base.update(overrides) return base -class TestClassifierReadLayers: +class TestProbeReadLayers: def test_omitted_reads_the_final_layer(self): - assert ( - GraniteSwitchConfig(**_classifier_kwargs()).classifier_read_layers is None - ) + assert GraniteSwitchConfig(**_probe_kwargs()).probe_read_layers is None @pytest.mark.parametrize("layers", [[None, None], [None, 0], [None, 3]]) def test_valid_layers_are_accepted(self, layers): - cfg = GraniteSwitchConfig(**_classifier_kwargs(classifier_read_layers=layers)) - assert cfg.classifier_read_layers == layers + cfg = GraniteSwitchConfig(**_probe_kwargs(probe_read_layers=layers)) + assert cfg.probe_read_layers == layers def test_wrong_length_raises(self): - with pytest.raises(ValueError, match="classifier_read_layers length"): - GraniteSwitchConfig( - **_classifier_kwargs(classifier_read_layers=[None, 2, 3]) - ) + with pytest.raises(ValueError, match="probe_read_layers length"): + GraniteSwitchConfig(**_probe_kwargs(probe_read_layers=[None, 2, 3])) - def test_non_classifier_slot_with_a_layer_raises(self): - with pytest.raises(ValueError, match="non-classifier slot 0"): - GraniteSwitchConfig(**_classifier_kwargs(classifier_read_layers=[1, 2])) + def test_non_probe_slot_with_a_layer_raises(self): + with pytest.raises(ValueError, match="non-probe slot 0"): + GraniteSwitchConfig(**_probe_kwargs(probe_read_layers=[1, 2])) @pytest.mark.parametrize("layer", [-1, 4, 99, True, 1.0]) def test_out_of_range_or_non_int_layer_raises(self, layer): # 4 == num_decoder_layers: the final state is requested by omitting it. with pytest.raises(ValueError, match=r"in 0\.\.3"): - GraniteSwitchConfig( - **_classifier_kwargs(classifier_read_layers=[None, layer]) - ) + GraniteSwitchConfig(**_probe_kwargs(probe_read_layers=[None, layer])) def test_dual_stream_with_intermediate_layer(self): cfg = GraniteSwitchConfig( - **_classifier_kwargs( - classifier_read_layers=[None, 2], + **_probe_kwargs( + probe_read_layers=[None, 2], dual_stream=True, cross_stream_rank=8, ) ) - assert cfg.classifier_read_layers == [None, 2] + assert cfg.probe_read_layers == [None, 2] restored = GraniteSwitchConfig.from_dict(cfg.to_dict()) assert restored.dual_stream - assert restored.classifier_read_layers == cfg.classifier_read_layers + assert restored.probe_read_layers == cfg.probe_read_layers - def test_classifier_slot_id_has_no_config_cap(self): - GraniteSwitchConfig(**_classifier_kwargs(num_adapters=257)) + def test_probe_slot_id_has_no_config_cap(self): + GraniteSwitchConfig(**_probe_kwargs(num_adapters=257)) def test_dual_stream_with_final_layer_is_accepted(self): - GraniteSwitchConfig(**_classifier_kwargs(dual_stream=True, cross_stream_rank=8)) + GraniteSwitchConfig(**_probe_kwargs(dual_stream=True, cross_stream_rank=8)) diff --git a/tests/unit/test_classifier_routing.py b/tests/unit/test_probe_routing.py similarity index 61% rename from tests/unit/test_classifier_routing.py rename to tests/unit/test_probe_routing.py index bdee44ed..35b74d5c 100644 --- a/tests/unit/test_classifier_routing.py +++ b/tests/unit/test_probe_routing.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -"""Classifier-slot routing: the kind LUT and the LoRA/classifier index split.""" +"""Probe-slot routing: the kind LUT and the LoRA/probe index split.""" import torch -from granite_switch.classifier_routing import ( +from granite_switch.config import GraniteSwitchConfig +from granite_switch.probe_routing import ( build_adapter_kind_lut, split_adapter_indices, ) -from granite_switch.config import GraniteSwitchConfig def _config(adapter_kinds, label_ids, base_reset=False): @@ -29,44 +29,44 @@ def _config(adapter_kinds, label_ids, base_reset=False): adapter_ranks=[4] * num_adapters, max_lora_rank=4, adapter_kinds=adapter_kinds, - classifier_label_token_ids=label_ids, + probe_label_token_ids=label_ids, ) class TestAdapterKindLut: - def test_none_without_classifier_slots(self): + def test_none_without_probe_slots(self): assert build_adapter_kind_lut(_config(["lora", "lora"], None)) is None - def test_marks_classifier_adapter_indices(self): - cfg = _config(["lora", "classifier", "classifier"], [None, [100], [101, 102]]) + def test_marks_probe_adapter_indices(self): + cfg = _config(["lora", "probe", "probe"], [None, [100], [101, 102]]) assert build_adapter_kind_lut(cfg).tolist() == [False, False, True, True] def test_base_reset_token_does_not_shift_slots(self): """The LUT is indexed by adapter index, not by control-token position.""" - cfg = _config(["lora", "classifier"], [None, [100, 101]], base_reset=True) + cfg = _config(["lora", "probe"], [None, [100, 101]], base_reset=True) assert build_adapter_kind_lut(cfg).tolist() == [False, False, True] - assert cfg.classifier_control_token_ids == [501] + assert cfg.probe_control_token_ids == [501] class TestSplitAdapterIndices: - def test_classifier_positions_leave_the_lora_stream(self): - cfg = _config(["lora", "classifier"], [None, [100, 101]]) - lora, classifier = split_adapter_indices( + def test_probe_positions_leave_the_lora_stream(self): + cfg = _config(["lora", "probe"], [None, [100, 101]]) + lora, probe = split_adapter_indices( build_adapter_kind_lut(cfg), torch.tensor([0, 1, 1, 0, 2]) ) assert lora.tolist() == [0, 1, 1, 0, 0] - assert classifier.tolist() == [0, 0, 0, 0, 2] + assert probe.tolist() == [0, 0, 0, 0, 2] def test_preserves_batched_shape(self): - cfg = _config(["classifier", "lora"], [[100, 101], None]) - lora, classifier = split_adapter_indices( + cfg = _config(["probe", "lora"], [[100, 101], None]) + lora, probe = split_adapter_indices( build_adapter_kind_lut(cfg), torch.tensor([[0, 2, 1], [1, 1, 0]]) ) assert lora.tolist() == [[0, 2, 0], [0, 0, 0]] - assert classifier.tolist() == [[0, 0, 1], [1, 1, 0]] + assert probe.tolist() == [[0, 0, 1], [1, 1, 0]] - def test_no_classifier_slots_pass_through(self): + def test_no_probe_slots_pass_through(self): indices = torch.tensor([0, 1, 2]) - lora, classifier = split_adapter_indices(None, indices) + lora, probe = split_adapter_indices(None, indices) assert lora is indices - assert not classifier.any() + assert not probe.any() diff --git a/tests/vllm/_model_forward_tests.py b/tests/vllm/_model_forward_tests.py index afc4ac1a..885e1bc2 100644 --- a/tests/vllm/_model_forward_tests.py +++ b/tests/vllm/_model_forward_tests.py @@ -61,7 +61,7 @@ def _try_import_vllm(): def _tiny_vllm_config( adapter_kinds=None, - classifier_label_token_ids=None, + probe_label_token_ids=None, ): """Minimal GraniteSwitchConfig for single-GPU vLLM tests.""" return GraniteSwitchConfig( @@ -80,7 +80,7 @@ def _tiny_vllm_config( max_lora_rank=16, adapter_ranks=[16, 16], adapter_kinds=adapter_kinds, - classifier_label_token_ids=classifier_label_token_ids, + probe_label_token_ids=probe_label_token_ids, switch_head_dim=32, max_position_embeddings=512, attention_multiplier=1.0, @@ -474,11 +474,11 @@ def test_different_adapters_produce_different_logits(self): ) -class TestClassifierSlotRouting(_VLLMModelTestBase): +class TestProbeSlotRouting(_VLLMModelTestBase): def _make_config(self): return _tiny_vllm_config( - adapter_kinds=["lora", "classifier"], - classifier_label_token_ids=[None, [100, 101]], + adapter_kinds=["lora", "probe"], + probe_label_token_ids=[None, [100, 101]], ) def test_marker_routes_out_of_lora_and_emits_label(self, monkeypatch): @@ -488,10 +488,10 @@ def test_marker_routes_out_of_lora_and_emits_label(self, monkeypatch): original_split = switch.split_indices def record_split(adapter_indices): - lora_indices, classifier_indices = original_split(adapter_indices) + lora_indices, probe_indices = original_split(adapter_indices) seen["lora"] = lora_indices.detach().cpu() - seen["classifier"] = classifier_indices.detach().cpu() - return lora_indices, classifier_indices + seen["probe"] = probe_indices.detach().cpu() + return lora_indices, probe_indices monkeypatch.setattr(switch, "split_indices", record_split) @@ -499,7 +499,7 @@ def record_split(adapter_indices): logits = self._run_forward_and_logits([10, 20, 30, 251]) assert seen["lora"].tolist() == [0, 0, 0, 0] - assert seen["classifier"].tolist() == [0, 0, 0, 2] + assert seen["probe"].tolist() == [0, 0, 0, 2] final_logits = logits[-1] assert torch.isfinite(final_logits[[100, 101]]).all() non_label_ids = [ diff --git a/tests/vllm/_classifier_chunked_prefill_worker.py b/tests/vllm/_probe_chunked_prefill_worker.py similarity index 91% rename from tests/vllm/_classifier_chunked_prefill_worker.py rename to tests/vllm/_probe_chunked_prefill_worker.py index c888c9d3..c80c2fff 100644 --- a/tests/vllm/_classifier_chunked_prefill_worker.py +++ b/tests/vllm/_probe_chunked_prefill_worker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Subprocess worker for the vLLM classifier chunked-prefill test. +"""Subprocess worker for the vLLM probe chunked-prefill test. Each sub-command is a separate process invocation so only one vLLM engine is ever resident (same discipline as ``_generation_equivalence_worker``):: @@ -36,7 +36,7 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): - """Compose a single-classifier-slot checkpoint. + """Compose a single-probe-slot checkpoint. ``num_labels`` single-token labels are taken from ``_LABEL_CANDIDATES`` and ``scale`` is the std of the random head weight. @@ -61,8 +61,8 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): adapter_token_ids=[control_id], adapter_substitute_token_ids=[control_id], adapter_names=["detect"], - adapter_kinds=["classifier"], - classifier_label_token_ids=[ids], + adapter_kinds=["probe"], + probe_label_token_ids=[ids], ) ckpt = os.path.join(outdir, "composed") model.save_pretrained(ckpt) @@ -80,7 +80,7 @@ def compose(base, outdir, control_offset=5, num_labels=6, scale=0.5): "owed to the moon each evening without fail or apparent complaint " "although sailors insisted otherwise over cheap wine and louder songs" ) - # The chat template's classifier layout: the user turn, its close, then the + # The chat template's probe layout: the user turn, its close, then the # marker as the very last token, which is itself the read point. body = tok.encode( f"<|start_of_role|>user<|end_of_role|>Is this text safe? {words}" @@ -151,15 +151,15 @@ def _fixed_head(seed, num_labels, hidden_size, scale=0.5): def _write_head(slot_dir, weight, bias, labels): from safetensors.torch import save_file - from granite_switch.composer.weight_transfer import CLASSIFIER_HEAD_FILE + from granite_switch.composer.weight_transfer import PROBE_HEAD_FILE os.makedirs(slot_dir, exist_ok=True) save_file( {"weight": weight.contiguous(), "bias": bias.contiguous()}, - os.path.join(slot_dir, CLASSIFIER_HEAD_FILE), + os.path.join(slot_dir, PROBE_HEAD_FILE), ) json.dump( - {"kind": "classifier", "labels": labels}, + {"kind": "probe", "labels": labels}, open(os.path.join(slot_dir, "adapter_config.json"), "w"), ) open(os.path.join(slot_dir, "io.yaml"), "w").write( @@ -172,7 +172,7 @@ def _prompt(tok, control_id): def compose_at_layer(base, outdir, read_layer, control_offset=5, num_labels=2): - """Compose one classifier slot reading ``read_layer``, plus its fp32 HF reference.""" + """Compose one probe slot reading ``read_layer``, plus its fp32 HF reference.""" import torch from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer @@ -199,9 +199,9 @@ def compose_at_layer(base, outdir, read_layer, control_offset=5, num_labels=2): adapter_token_ids=[control_id], adapter_substitute_token_ids=[control_id], adapter_names=["detect"], - adapter_kinds=["classifier"], - classifier_label_token_ids=[ids], - classifier_read_layers=[read_layer], + adapter_kinds=["probe"], + probe_label_token_ids=[ids], + probe_read_layers=[read_layer], ) ckpt = os.path.join(outdir, "composed") model.save_pretrained(ckpt) @@ -245,9 +245,9 @@ def compose_two_slot(base, outdir, read_layer, control_offset=5, num_labels=2): adapter_token_ids=control_ids, adapter_substitute_token_ids=control_ids, adapter_names=["detect_final", "detect_mid"], - adapter_kinds=["classifier", "classifier"], - classifier_label_token_ids=[ids, ids], - classifier_read_layers=[None, read_layer], + adapter_kinds=["probe", "probe"], + probe_label_token_ids=[ids, ids], + probe_read_layers=[None, read_layer], ) ckpt = os.path.join(outdir, "composed_two_slot") model.save_pretrained(ckpt) @@ -269,25 +269,25 @@ def compose_two_slot(base, outdir, read_layer, control_offset=5, num_labels=2): return result -def _record_classifier_rows(llm): +def _record_probe_rows(llm): """Install after startup; PP callbacks must be importable for standard pickle.""" - from tests.shared.classifier_recorder import install_classifier_recorder + from tests.shared.probe_recorder import install_probe_recorder - llm.apply_model(install_classifier_recorder) + llm.apply_model(install_probe_recorder) -def _read_classifier_rows(llm): - from tests.shared.classifier_recorder import read_classifier_records +def _read_probe_rows(llm): + from tests.shared.probe_recorder import read_probe_records - return [row for rows in llm.apply_model(read_classifier_records) for row in rows] + return [row for rows in llm.apply_model(read_probe_records) for row in rows] def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, eager): - """Serve both classifier slots' prompts, one request each, eager or FULL. + """Serve both probe slots' prompts, one request each, eager or FULL. Under FULL both requests share one captured batch shape, exercising replay with two different slots/read layers. Reports per-label logprobs - and the recorded classifier row (packed scores, rewritten label logits, + and the recorded probe row (packed scores, rewritten label logits, and for the final-layer slot its actual activation and loaded head). """ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" @@ -305,7 +305,7 @@ def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, if not eager: kw["compilation_config"] = {"cudagraph_mode": "FULL"} llm = LLM(**kw) - _record_classifier_rows(llm) + _record_probe_rows(llm) sp = SamplingParams(max_tokens=1, temperature=0.0, logprobs=len(label_token_ids)) results = [] @@ -313,7 +313,7 @@ def run_two_slot(ckpt, prompt_ids_1, prompt_ids_2, label_token_ids, out_path, *, out = llm.generate([TokensPrompt(prompt_token_ids=prompt_ids)], sp)[0].outputs[ 0 ] - records = _read_classifier_rows(llm) + records = _read_probe_rows(llm) recorded_slots = [r["slot"] for r in records] if recorded_slots != [expected_slot]: raise RuntimeError( @@ -369,7 +369,7 @@ def run_at_layer(ckpt, prompt_ids, label_token_ids, out_path, full_cudagraph=Fal ) -def run_sr_classifier(ckpt, out_path, *, eager, pp_size=1): +def run_sr_probe(ckpt, out_path, *, eager, pp_size=1): """Serve each SR-checkpoint probe and an SR request, sequentially and mixed. Each slot is selected on its own: a probe request carries only its marker, @@ -388,7 +388,7 @@ def run_sr_classifier(ckpt, out_path, *, eager, pp_size=1): labels = list(range(100, 100 + HIDDEN)) config = AutoConfig.from_pretrained(ckpt) - markers = config.classifier_control_token_ids + markers = config.probe_control_token_ids kw = dict( model=ckpt, skip_tokenizer_init=True, @@ -410,7 +410,7 @@ def run_sr_classifier(ckpt, out_path, *, eager, pp_size=1): "cudagraph_capture_sizes": [1, 2, 4, 16], } llm = LLM(**kw) - _record_classifier_rows(llm) + _record_probe_rows(llm) sampling = SamplingParams(max_tokens=1, temperature=0.0, logprobs=HIDDEN + 1) sr_request = [10, CONTROL_TOKEN_ID, 20, 30] @@ -419,7 +419,7 @@ def serve(prompts): [TokensPrompt(prompt_token_ids=p) for p in prompts], sampling ) return { - "records": _read_classifier_rows(llm), + "records": _read_probe_rows(llm), "outputs": [ { "token_id": int(o.outputs[0].token_ids[0]), @@ -501,7 +501,7 @@ def summarize(output): o = llm.generate([prompt], sp)[0].outputs[0] results.append(summarize(o)) - # Same shape (base_len + 1 tokens), no classifier marker at all -- the + # Same shape (base_len + 1 tokens), no probe marker at all -- the # verdict from the PRECEDING call must not leak into this one. This is a # SEQUENTIAL check (separate generate() calls), not a same-batch mix. plain_body = [rng.randint(1000, 2000) for _ in range(base_len + 1)] @@ -509,11 +509,11 @@ def summarize(output): o = llm.generate([plain_prompt], sp)[0].outputs[0] plain_result = summarize(o) - # A GENUINE mixed batch: 3 requests in ONE generate() call -- classifier, - # plain, classifier -- so vLLM schedules and samples all three rows in the + # A GENUINE mixed batch: 3 requests in ONE generate() call -- probe, + # plain, probe -- so vLLM schedules and samples all three rows in the # same forward pass. Each must get its own correct result: the plain - # row's logits must not pick up either classifier row's verdict, and the - # two classifier rows (same slot here, since this checkpoint has one) + # row's logits must not pick up either probe row's verdict, and the + # two probe rows (same slot here, since this checkpoint has one) # must each still fire independently. mix_cls_body_1 = [rng.randint(1000, 2000) for _ in range(base_len)] mix_plain_body = [rng.randint(1000, 2000) for _ in range(base_len + 1)] @@ -569,21 +569,19 @@ def run_budget(ckpt, prompt_ids, budget, out_path): def main(): mode = sys.argv[1] - if mode in ("compose-sr-classifier", "compose-sr-classifiers"): + if mode in ("compose-sr-probe", "compose-sr-probes"): from pathlib import Path - from tests.composer.test_sr_compose_e2e import _compose_sr_classifier + from tests.composer.test_sr_compose_e2e import _compose_sr_probe - read_layers = [0, 1, None] if mode == "compose-sr-classifiers" else [None] - model = _compose_sr_classifier( - Path(os.environ["CLS_OUT"]), read_layers=read_layers - ) + read_layers = [0, 1, None] if mode == "compose-sr-probes" else [None] + model = _compose_sr_probe(Path(os.environ["CLS_OUT"]), read_layers=read_layers) ckpt = os.path.join(os.environ["CLS_OUT"], "composed") model.save_pretrained(ckpt) print(json.dumps({"result": {"ckpt": ckpt}})) return if mode in ("run-sr-eager", "run-sr-full", "run-sr-pp-eager", "run-sr-pp-full"): - run_sr_classifier( + run_sr_probe( os.environ["CLS_CKPT"], sys.argv[3] if mode.startswith("run-sr-pp-") else sys.argv[2], eager=mode.endswith("eager"), diff --git a/tests/vllm/test_classifier_capture_safety.py b/tests/vllm/test_probe_capture_safety.py similarity index 80% rename from tests/vllm/test_classifier_capture_safety.py rename to tests/vllm/test_probe_capture_safety.py index 31f1c32b..e9fe64df 100644 --- a/tests/vllm/test_classifier_capture_safety.py +++ b/tests/vllm/test_probe_capture_safety.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -"""CPU method tests for token-aligned, score-at-layer classification. +"""CPU method tests for token-aligned, score-at-layer probe verdicts. GPU capture/replay and packed-output transport are tested in -``test_classifier_chunked_prefill.py``. Request-boundary validation is deferred. +``test_probe_chunked_prefill.py``. Request-boundary validation is deferred. """ from types import SimpleNamespace @@ -24,11 +24,11 @@ ) -class _StubClassifierHead: +class _StubProbeHead: """Score = slot id * hidden[:, 0], so a score shows which slot and source it read.""" - def __call__(self, x, classifier_indices): - return (classifier_indices.to(x.dtype) * x[:, 0]).unsqueeze(-1).expand(-1, 2) + def __call__(self, x, probe_indices): + return (probe_indices.to(x.dtype) * x[:, 0]).unsqueeze(-1).expand(-1, 2) class _Model(torch.nn.Module): @@ -43,32 +43,30 @@ def __init__(self, slot_of_marker, read_layers=None, kinds=None): super().__init__() num_slots = max(slot_of_marker.values()) read_layers = read_layers or [None] * num_slots - is_classifier = [k == "classifier" for k in kinds or ["classifier"] * num_slots] - self.config = type("Cfg", (), {"max_classifier_labels": 2})() - self.classifier_head = _StubClassifierHead() - self.classifier_slot_mask_t = torch.tensor([False, *is_classifier]) + is_probe = [k == "probe" for k in kinds or ["probe"] * num_slots] + self.config = type("Cfg", (), {"max_probe_labels": 2})() + self.probe_head = _StubProbeHead() + self.probe_slot_mask_t = torch.tensor([False, *is_probe]) slots_by_layer = {} - for slot, (is_clf, layer) in enumerate( - zip(is_classifier, read_layers), start=1 - ): + for slot, (is_clf, layer) in enumerate(zip(is_probe, read_layers), start=1): if is_clf: slots_by_layer.setdefault(layer, []).append(slot) - self._classifier_final_layer_group = None - self._classifier_layer_groups = {} + self._probe_final_layer_group = None + self._probe_layer_groups = {} for layer, slot_ids in slots_by_layer.items(): - name = f"_classifier_slots_{'final' if layer is None else layer}" + name = f"_probe_slots_{'final' if layer is None else layer}" self.register_buffer(name, torch.tensor(slot_ids)) if layer is None: - self._classifier_final_layer_group = name + self._probe_final_layer_group = name else: - self._classifier_layer_groups[layer] = name + self._probe_layer_groups[layer] = name def _run(model, slot_of_marker, tokens): - """Call ``_compute_classifier_outputs`` on a flat token list; hidden[t] = t + 1.""" + """Call ``_compute_probe_outputs`` on a flat token list; hidden[t] = t + 1.""" indices = [slot_of_marker.get(t, 0) for t in tokens] hidden = torch.arange(1.0, len(tokens) + 1).unsqueeze(1).expand(-1, 4) - return GraniteSwitchModel._compute_classifier_outputs( + return GraniteSwitchModel._compute_probe_outputs( model, hidden, torch.tensor(indices), torch.zeros(len(tokens), 2) ) @@ -77,7 +75,7 @@ def _score(model, layer, hidden_value, indices, scores): hidden = torch.full((len(indices), 4), hidden_value) return GraniteSwitchModel._score_group( model, - model._classifier_layer_groups[layer], + model._probe_layer_groups[layer], hidden, torch.tensor(indices), scores, @@ -113,7 +111,7 @@ def test_no_marker_is_all_zero(): def test_transferred_marker_indices_need_no_input_ids(): model = _Model({90: 2, 91: 3}) - scores, slot = GraniteSwitchModel._compute_classifier_outputs( + scores, slot = GraniteSwitchModel._compute_probe_outputs( model, torch.ones(3, 4), torch.tensor([0, 2, 0]), torch.zeros(3, 2) ) assert slot.tolist() == [0, 2, 0] @@ -123,7 +121,7 @@ def test_transferred_marker_indices_need_no_input_ids(): def test_final_layer_scoring_leaves_intermediate_slot_scores_untouched(): model = _Model({90: 2, 91: 3}, read_layers=[None, 1, None]) pre_scores = torch.tensor([[0.0, 0.0], [99.0, 99.0], [0.0, 0.0]]) - scores, slot = GraniteSwitchModel._compute_classifier_outputs( + scores, slot = GraniteSwitchModel._compute_probe_outputs( model, torch.ones(3, 4), torch.tensor([0, 2, 3]), @@ -146,8 +144,8 @@ def test_each_intermediate_layer_scores_only_its_own_slots(): def test_slots_sharing_a_layer_form_one_group(): model = _Model({90: 1, 91: 2}) - assert model._classifier_layer_groups == {} - assert getattr(model, model._classifier_final_layer_group).tolist() == [1, 2] + assert model._probe_layer_groups == {} + assert getattr(model, model._probe_final_layer_group).tolist() == [1, 2] scores, slot = _run(model, {90: 1, 91: 2}, [90, 5, 91]) assert slot.tolist() == [1, 0, 2] assert scores[0].tolist() == [1.0, 1.0] @@ -155,9 +153,9 @@ def test_slots_sharing_a_layer_form_one_group(): def test_lora_slots_form_no_group(): - model = _Model({91: 2}, kinds=["lora", "classifier"]) - assert getattr(model, model._classifier_final_layer_group).tolist() == [2] - assert model._classifier_layer_groups == {} + model = _Model({91: 2}, kinds=["lora", "probe"]) + assert getattr(model, model._probe_final_layer_group).tolist() == [2] + assert model._probe_layer_groups == {} _, slot = _run(model, {91: 2}, [5, 91, 6]) assert slot.tolist() == [0, 2, 0] @@ -166,13 +164,11 @@ def test_profiling_random_slot_on_a_lora_slot_leaves_logits_untouched(): """vLLM's profiling sampler run feeds random hidden states; 0.75 rounds to LoRA slot 1.""" wrapper = GraniteSwitchForCausalLM.__new__(GraniteSwitchForCausalLM) torch.nn.Module.__init__(wrapper) - wrapper.model = _Model({91: 2}, kinds=["lora", "classifier"]) - wrapper.config = type( - "Cfg", (), {"classifier_label_token_ids": [None, [100, 200]]} - )() + wrapper.model = _Model({91: 2}, kinds=["lora", "probe"]) + wrapper.config = type("Cfg", (), {"probe_label_token_ids": [None, [100, 200]]})() logits = torch.randn(1, 300) expected = logits.clone() - out = GraniteSwitchForCausalLM._apply_classifier_verdict( + out = GraniteSwitchForCausalLM._apply_probe_verdict( wrapper, logits, torch.randn(1, 2), torch.tensor([0.75]) ) assert torch.equal(out, expected) @@ -186,7 +182,7 @@ def test_scoring_resolves_replaced_registered_buffers(): for tensor in old_buffers: tensor.zero_() scores = _score(model, 0, 3.0, [1, 2], torch.zeros(2, 2)) - scores, slots = GraniteSwitchModel._compute_classifier_outputs( + scores, slots = GraniteSwitchModel._compute_probe_outputs( model, torch.full((2, 4), 5.0), torch.tensor([1, 2]), @@ -206,7 +202,7 @@ def test_pipeline_handoff_without_later_token_ids(monkeypatch, sr, stages, dtype on the last rank. With three stages, the middle rank only relays scores. Distinct values expose resets, wrong layers, stream selection and reliance on input_ids. The switch keeps the latest slot, like MultiSwitch, so only - the markers themselves may carry a classifier index. Markers route out of + the markers themselves may carry a probe index. Markers route out of the LoRA stream; the trailing LoRA token reads the adapted SR stream. The token slice matches vLLM's PP buffer contract. """ @@ -222,11 +218,11 @@ def __call__(self, input_ids, **kwargs): return torch.tensor(indices), input_ids def split_indices(self, indices): - is_classifier = indices >= 2 + is_probe = indices >= 2 zeros = torch.zeros_like(indices) return ( - torch.where(is_classifier, zeros, indices), - torch.where(is_classifier, indices, zeros), + torch.where(is_probe, zeros, indices), + torch.where(is_probe, indices, zeros), ) class Layer(torch.nn.Module): @@ -241,20 +237,20 @@ def forward(self, positions, hidden_states, residual=None): model = _Model( {90: 2, 91: 3, 92: 4}, read_layers=[None, 0, stages - 1, None], - kinds=["lora", "classifier", "classifier", "classifier"], + kinds=["lora", "probe", "probe", "probe"], ) model.config.hidden_size = 4 model.config.embedding_multiplier = 1.0 model.config.fused_add_norm = False model.switch = Switch() model.adapter_token_ids = torch.tensor([89, 90, 91, 92]) - model.classifier_control_token_ids_t = torch.tensor([90, 91, 92]) + model.probe_control_token_ids_t = torch.tensor([90, 91, 92]) model.lora_meta = model.lora_ctx = None model.decoder_interface = SRDecoderInterface() if sr else LoRADecoderInterface() model.layers = torch.nn.ModuleList([Layer() for _ in range(stages)]) model.norm = torch.nn.Identity() - model._compute_classifier_outputs = ( - GraniteSwitchModel._compute_classifier_outputs.__get__(model) + model._compute_probe_outputs = GraniteSwitchModel._compute_probe_outputs.__get__( + model ) model.start_layer, model.end_layer = 0, 1 monkeypatch.setattr( @@ -269,8 +265,8 @@ def forward(self, positions, hidden_states, residual=None): inputs_embeds=torch.ones(5, 4, dtype=dtype), )[1:] assert isinstance(wire, IntermediateTensors) - assert wire["classifier_indices"].tolist() == [2, 3, 4, 0] - assert wire["classifier_scores"][:, 0].tolist() == [4, 0, 0, 0] + assert wire["probe_indices"].tolist() == [2, 3, 4, 0] + assert wire["probe_scores"][:, 0].tolist() == [4, 0, 0, 0] buffers = GraniteSwitchModel.make_empty_intermediate_tensors( model, 4, dtype, torch.device("cpu") ) @@ -288,7 +284,7 @@ def forward(self, positions, hidden_states, residual=None): model, None, torch.arange(4), intermediate_tensors=buffers ) if not group.is_last_rank: - assert torch.equal(packed["classifier_scores"], wire["classifier_scores"]) + assert torch.equal(packed["probe_scores"], wire["probe_scores"]) for key, tensor in packed.items(): buffers[key].copy_(tensor) assert packed[:, -1].tolist() == [2, 3, 4, 0] diff --git a/tests/vllm/test_classifier_chunked_prefill.py b/tests/vllm/test_probe_chunked_prefill.py similarity index 94% rename from tests/vllm/test_classifier_chunked_prefill.py rename to tests/vllm/test_probe_chunked_prefill.py index 09fb2e62..53a54f96 100644 --- a/tests/vllm/test_classifier_chunked_prefill.py +++ b/tests/vllm/test_probe_chunked_prefill.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -"""GPU coverage for the vLLM classifier verdict exit: FULL capture/replay, +"""GPU coverage for the vLLM probe verdict exit: FULL capture/replay, chunked prefill, intermediate-layer reads, and two slots in one checkpoint. ``enable_chunked_prefill`` is on by default in serving. The prompt is laid out as -the chat template renders it, with the classifier control token as the very last +the chat template renders it, with the probe control token as the very last token, so it always lands in the final chunk. The same prompt served at ``max_num_batched_tokens`` in {unchunked, 256, 101, 48} must emit the same label token. @@ -26,12 +26,10 @@ import torch pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="vLLM classifier serving test requires a GPU" + not torch.cuda.is_available(), reason="vLLM probe serving test requires a GPU" ) -WORKER = os.path.join( - os.path.dirname(__file__), "_classifier_chunked_prefill_worker.py" -) +WORKER = os.path.join(os.path.dirname(__file__), "_probe_chunked_prefill_worker.py") _REPO = os.path.dirname(os.path.dirname(os.path.dirname(WORKER))) BASE_MODEL = os.environ.get("CLS_TEST_BASE", "ibm-granite/granite-4.1-3b") @@ -50,7 +48,7 @@ def _env(): def _compose(tmp_path_factory, argv, tag): - """Compose a classifier checkpoint in a CPU subprocess, before any CUDA init.""" + """Compose a probe checkpoint in a CPU subprocess, before any CUDA init.""" outdir = str(tmp_path_factory.mktemp(tag)) proc = subprocess.run( [sys.executable, WORKER, *argv], @@ -110,7 +108,7 @@ def _run_worker(outdir, argv, env, timeout=1800): @pytest.fixture(scope="module") def composed(tmp_path_factory): - """Multi-label classifier over a varied prompt.""" + """Multi-label probe over a varied prompt.""" result = _compose(tmp_path_factory, ["compose-multilabel"], "cls_ckpt") result["pids_path"] = _write_pids(result) return result @@ -122,16 +120,16 @@ def _composed_env(composed, **extra): @pytest.fixture(scope="module") def composed_sr(tmp_path_factory): - return _compose(tmp_path_factory, ["compose-sr-classifier"], "sr_cls_ckpt") + return _compose(tmp_path_factory, ["compose-sr-probe"], "sr_cls_ckpt") @pytest.fixture(scope="module") def composed_sr_probes(tmp_path_factory): - return _compose(tmp_path_factory, ["compose-sr-classifiers"], "sr_probes_ckpt") + return _compose(tmp_path_factory, ["compose-sr-probes"], "sr_probes_ckpt") @pytest.mark.parametrize("mode", ["eager", "full"]) -def test_sr_checkpoint_serves_final_classifier_slot(composed_sr, mode): +def test_sr_checkpoint_serves_final_probe_slot(composed_sr, mode): """The identity head scores the served hidden state; transport stays exact.""" served = _run_worker( composed_sr["outdir"], @@ -209,7 +207,7 @@ def _check_sr_layer_probes(result, mode, pp): assert not failures, "\n".join(failures) -def test_sr_classifier_reads_each_layer_under_replay(sr_probe_reference): +def test_sr_probe_reads_each_layer_under_replay(sr_probe_reference): """One GPU: probes at layer 0, layer 1 and final norm beside an SR slot.""" mode, served = sr_probe_reference _check_sr_layer_probes(served, mode, 1) @@ -258,7 +256,7 @@ def test_full_cudagraph_mode_does_not_crash_and_matches_eager(composed): for i, (f, e) in enumerate(zip(full["results"], eager["results"])): assert f["is_label"], ( f"call {i}: FULL-cudagraph verdict {f['token_id']} ({f['text']!r}) " - f"is not a label token {label_ids}; the classifier exit did not " + f"is not a label token {label_ids}; the probe exit did not " f"fire under real CUDA-graph capture/replay" ) assert f["token_id"] == e["token_id"], ( @@ -279,7 +277,7 @@ def test_full_cudagraph_mode_does_not_crash_and_matches_eager(composed): assert actual["is_label"] -def test_classifier_token_stable_across_chunk_budgets(composed): +def test_probe_token_stable_across_chunk_budgets(composed): """Emitted verdict token identical across chunk budgets.""" label_ids = set(composed["label_token_ids"]) assert len(composed["prompt_ids"]) > min(b for b in BUDGETS if b), ( @@ -299,7 +297,7 @@ def test_classifier_token_stable_across_chunk_budgets(composed): ref = results[None] assert ref["token_id"] in label_ids, ( f"unchunked verdict {ref['token_id']} is not a label token {label_ids}; " - f"the classifier exit did not fire" + f"the probe exit did not fire" ) for b, r in results.items(): assert r["token_id"] == ref["token_id"], ( @@ -309,7 +307,7 @@ def test_classifier_token_stable_across_chunk_budgets(composed): # Score-gap tolerance against the fp32 HF reference; this comparison also -# includes HF/vLLM backbone differences, not only classifier arithmetic. +# includes HF/vLLM backbone differences, not only probe arithmetic. SCORE_DIFF_ATOL = 0.3 @@ -449,7 +447,7 @@ def gap(scores): print(f" activation contribution: {vllm_gap - hf_gap:.4f}") -def test_classifier_reads_an_intermediate_layer_correctly(tmp_path_factory): +def test_probe_reads_an_intermediate_layer_correctly(tmp_path_factory): """A layer-20 slot under eager and FULL, against an fp32 HF layer-20 reference. The fixed seed-0 head's reference margin (about 0.07) is below the score @@ -485,10 +483,10 @@ def test_classifier_reads_an_intermediate_layer_correctly(tmp_path_factory): assert not failures, "\n".join(failures) -def test_two_classifier_slots_replay_correctly_in_one_captured_shape(tmp_path_factory): +def test_two_probe_slots_replay_correctly_in_one_captured_shape(tmp_path_factory): """Serve a final-layer and a layer-15 slot under eager and FULL. - Classifier arithmetic and transport are checked against the head's + Probe arithmetic and transport are checked against the head's actual input, recorded in-process at ``compute_logits``: - slot 1 (final layer): packed scores equal the loaded bf16 head applied diff --git a/tutorials/README.md b/tutorials/README.md index 969b4541..f745b951 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -26,7 +26,7 @@ Step-by-step walkthroughs covering adapter function invocation, pipeline constru |-------|-------------| | [Using Mellea with Granite Switch](guides/mellea_with_granite_switch.md) | Connect Mellea to a Granite Switch model | | [Bring Your Own Adapter](guides/build_your_own_adapter.md) | Train, compose, and use custom adapters | -| [Classifier Slots](guides/classifier_slots.md) | Compose and serve a classifier head that emits a label word | +| [Activation Probe Slots](guides/activation_probe_slots.md) | Compose and serve an activation probe head that emits a label word | | [Compare Inference Throughput](guides/compare_inference_throughput.md) | Compare LoRA vs aLoRA based models in an inference race setup | diff --git a/tutorials/guides/classifier_slots.md b/tutorials/guides/activation_probe_slots.md similarity index 75% rename from tutorials/guides/classifier_slots.md rename to tutorials/guides/activation_probe_slots.md index 291c68df..3928a1d6 100644 --- a/tutorials/guides/classifier_slots.md +++ b/tutorials/guides/activation_probe_slots.md @@ -1,15 +1,16 @@ -# Classifier Slots +# Activation Probe Slots -This guide explains how to compose a Granite Switch model with a **classifier slot** — a -lightweight classification head fired by the same control token as a LoRA adapter — and how to -serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detection request. +This guide explains how to compose a Granite Switch model with an **activation probe slot** — a +lightweight linear head over a hidden state, fired by the same control token as a LoRA adapter — +and how to serve it so the model emits a label *word* (`safe`, `unsafe`, …) for a detection +request. ## Overview -A **classifier slot** holds a small linear classification head instead of a LoRA. It uses the same -control-token machinery (one `<|name|>` token, one adapter index) and reads the configured layer's -activation to produce per-label logits. Classifier and generation-adapter slots can coexist in one -checkpoint; a request selects one slot. +An **activation probe slot** (probe for short) holds a small linear head instead of a LoRA. It +uses the same control-token machinery (one `<|name|>` token, one adapter index) and reads the +configured layer's activation to produce per-label logits. Probe and generation-adapter slots can +coexist in one checkpoint; a request selects one slot. At inference the verdict does not come back as a raw score vector. Each label word is resolved to a single token id at compose time, and the model rewrites the request's final logit row so that the @@ -22,21 +23,21 @@ custom client. (weight/bias) (io.yaml or YAML list) (granite-switch) (vLLM OpenAI API) ``` -Multiple classifier slots may coexist, and **each slot may have its own number of labels** (e.g. a +Multiple probe slots may coexist, and **each slot may have its own number of labels** (e.g. a 2-label `safe`/`unsafe` slot next to a 3-label `good`/`bad`/`neutral` slot). ## Prerequisites - Granite base model (e.g., `ibm-granite/granite-4.1-3b`) -- A trained classifier head (see Step 1) +- A trained probe head (see Step 1) - `granite-switch[vllm,compose]` installed See [PREREQUISITES.md](../PREREQUISITES.md) for detailed setup. -## Step 1: Train (or stage) a classifier head +## Step 1: Train (or stage) a probe head -A classifier slot is **not** a LoRA, so its directory holds a different artifact. There is no -`adapter_model.safetensors`; instead the directory must contain a **`classifier_head.safetensors`** +A probe slot is **not** a LoRA, so its directory holds a different artifact. There is no +`adapter_model.safetensors`; instead the directory must contain a **`probe_head.safetensors`** with exactly two tensors: | Key | Shape | Meaning | @@ -71,7 +72,7 @@ from safetensors.torch import save_file HIDDEN = 2560 # granite-4.1-3b hidden size def make_head(name, labels): - """Write a synthetic classifier head for `name` under the library layout. + """Write a synthetic probe head for `name` under the library layout. A real head would carry trained weights; this one is zero-weight with a bias that forces row 0 to win, so the emitted label is deterministic. @@ -81,8 +82,8 @@ def make_head(name, labels): weight = torch.zeros(len(labels), HIDDEN) bias = torch.full((len(labels),), -10.0) bias[0] = 10.0 # row 0 (labels[0]) always wins - save_file({"weight": weight, "bias": bias}, str(head_dir / "classifier_head.safetensors")) - # Every composed slot (LoRA or classifier) also needs an io.yaml — compose copies + save_file({"weight": weight, "bias": bias}, str(head_dir / "probe_head.safetensors")) + # Every composed slot (LoRA or probe) also needs an io.yaml — compose copies # it verbatim into the checkpoint's io_configs/. A two-line stub is enough here. (head_dir / "io.yaml").write_text(f"name: {name}\nmodel: ~\n") return str(head_dir) @@ -90,7 +91,7 @@ def make_head(name, labels): > The directory layout `////` mirrors the library layout the > composer expects for a local adapter. The trailing `lora`/`alora` segment selects **token -> placement** (LoRA vs activated-LoRA), independent of the fact that this slot is a classifier. +> placement** (LoRA vs activated-LoRA), independent of the fact that this slot is a probe. > > The `io.yaml` is required for **every** slot: compose copies each slot's `io.yaml` into the > output's `io_configs//`, so a slot without one fails the copy step. Real adapters ship a @@ -104,7 +105,7 @@ token id per label. If a label word does not encode to a single token in the bas compose fails loud and lists the offenders, e.g.: ``` -1 classifier label word(s) do not encode to a single token, ... +1 probe label word(s) do not encode to a single token, ... 'Ambiguous' → 3 tokens: ['Amb', 'igu', 'ous'] ``` @@ -114,24 +115,24 @@ If no natural single-token word fits, the tokenizer's reserved `<|unused_N|>` to ## Step 2: Declare the slot in a manifest -Classifier slots are declared with two extra keys — `kind: classifier` and `labels:` — on top of the +Probe slots are declared with two extra keys — `kind: probe` and `labels:` — on top of the normal adapter fields. The simplest way to pass one or more of them to compose is a **YAML manifest** that maps each slot name to its path and metadata: ```yaml -# ./adapters.yaml (one classifier slot) +# ./adapters.yaml (one probe slot) safety: path: ./safety-head/safety/granite-4.1-3b/lora type: lora # token placement (lora vs alora); independent of `kind` - kind: classifier # marks this slot a classifier head (default "lora" when omitted) + kind: probe # marks this slot a probe head (default "lora" when omitted) labels: [safe, unsafe] # label words, resolved to single token ids at compose ``` `type` and `kind` are independent axes: `type` is the LoRA *technology* (drives control-token -placement), while `kind` is whether the slot is a real LoRA or a classifier head. A plain LoRA entry +placement), while `kind` is whether the slot is a real LoRA or a probe head. A plain LoRA entry simply omits `kind`/`labels`. -A classifier entry may also set an optional `layer`, the zero-based decoder layer whose output the +A probe entry may also set an optional `layer`, the zero-based decoder layer whose output the head reads (the residual stream after that layer, before the final norm). Omit it to read the final post-norm hidden state: @@ -139,42 +140,42 @@ post-norm hidden state: safety_mid: path: ./safety-mid-head/safety_mid/granite-4.1-3b/lora type: lora - kind: classifier + kind: probe labels: [safe, unsafe] layer: 20 # 0..num_decoder_layers - 1; omit for final ``` -Classifier slots can also be composed alongside Shadow Residual (SR) adapters; a classifier +Probe slots can also be composed alongside Shadow Residual (SR) adapters; a probe request reads that layer's base stream. Decoder layer numbers exclude the switch's KV-cache slots. Pipeline parallelism preserves the same numbering across stages. -You can mix LoRA and classifier entries in one manifest, and give each classifier its own `labels` +You can mix LoRA and probe entries in one manifest, and give each probe its own `labels` list of any length: ```yaml -# ./adapters.yaml — one real LoRA + two synthetic classifier slots (2 and 3 labels) +# ./adapters.yaml — one real LoRA + two synthetic probe slots (2 and 3 labels) answerability: path: /path/to/answerability/granite-4.1-3b/lora # a real LoRA (see below) type: lora safety: path: ./safety-head/safety/granite-4.1-3b/lora type: lora - kind: classifier + kind: probe labels: [safe, unsafe] faithfulness: path: ./faithfulness-head/faithfulness/granite-4.1-3b/lora type: lora - kind: classifier + kind: probe labels: [good, bad, neutral] # single-token words — see the label rule above ``` > **Slot ordering is automatic.** Compose stably reorders slots so all LoRA slots come first and -> classifier slots last, then derives every downstream artifact (control tokens, chat template, -> adapter indices) from that order. You do not need to order the manifest yourself; classifiers +> probe slots last, then derives every downstream artifact (control tokens, chat template, +> adapter indices) from that order. You do not need to order the manifest yourself; probes > ending up as the contiguous suffix is what lets the LoRA weight transfer stay positional and the -> classifier head bank pad uniformly. +> probe head bank pad uniformly. -The script below download a rag lora and constructed two synthetic classifiers for testing. +The script below download a rag lora and constructed two synthetic probes for testing. ```python from pathlib import Path @@ -189,16 +190,16 @@ snap = snapshot_download( ) lora_dir = str(Path(snap) / "answerability" / "granite-4.1-3b" / "lora") -# 2. Two synthetic classifier slots with different label counts (make_head from Step 1). +# 2. Two synthetic probe slots with different label counts (make_head from Step 1). safety_dir = make_head("safety", ["safe", "unsafe"]) faith_dir = make_head("faithfulness", ["good", "bad", "neutral"]) -# 3. Emit the manifest pairing the real LoRA with the two synthetic classifiers. +# 3. Emit the manifest pairing the real LoRA with the two synthetic probes. manifest = { "answerability": {"path": lora_dir, "type": "lora"}, - "safety": {"path": safety_dir, "type": "lora", "kind": "classifier", + "safety": {"path": safety_dir, "type": "lora", "kind": "probe", "labels": ["safe", "unsafe"]}, - "faithfulness": {"path": faith_dir, "type": "lora", "kind": "classifier", + "faithfulness": {"path": faith_dir, "type": "lora", "kind": "probe", "labels": ["good", "bad", "neutral"]}, } Path("./adapters.yaml").write_text(yaml.safe_dump(manifest, sort_keys=False)) @@ -221,37 +222,37 @@ python -m granite_switch.composer.compose_granite_switch \ Compose will: 1. Load the base model and discover the slots from the manifest. -2. Resolve each classifier's label words to single token ids (failing loud on multi-token words). -3. Read each classifier's `classifier_head.safetensors` into the head bank at its slot. +2. Resolve each probe's label words to single token ids (failing loud on multi-token words). +3. Read each probe's `probe_head.safetensors` into the head bank at its slot. 4. Stack everything into one checkpoint under `--output`. ``` -Classifier slots: +Probe slots: safety → labels ['safe', 'unsafe'] → token ids [19193, 39257] @ final norm faithfulness → labels ['good', 'bad', 'neutral'] → token ids [19045, 14176, 60668] @ final norm ``` A slot with `layer: 20` prints `@ layer 20` instead. -The resulting `config.json` carries the per-slot classifier metadata as plain lists/ints, so it +The resulting `config.json` carries the per-slot probe metadata as plain lists/ints, so it round-trips through `save_pretrained`/`from_pretrained` with no custom code, and both the HF and vLLM backends read it identically. Three fields are the stored contract: | Field | Meaning | |-------|---------| -| `adapter_kinds` | per-slot `"lora"` / `"classifier"`, one entry per adapter — what marks a slot a classifier | -| `classifier_label_token_ids` | per-slot list of label token ids, `null` on LoRA slots | -| `classifier_read_layers` | per-slot zero-based decoder layer a classifier reads; `null` for the final post-norm state and on LoRA slots | +| `adapter_kinds` | per-slot `"lora"` / `"probe"`, one entry per adapter — what marks a slot a probe | +| `probe_label_token_ids` | per-slot list of label token ids, `null` on LoRA slots | +| `probe_read_layers` | per-slot zero-based decoder layer a probe reads; `null` for the final post-norm state and on LoRA slots | The rest are **derived** from those in `GraniteSwitchConfig.__init__`, so they appear on the config object but are not independent inputs: | Derived field | From | |---------------|------| -| `classifier_num_labels_per_slot` | `len()` of each slot's label token ids (`0` on LoRA slots) | -| `max_classifier_labels` | `max()` of the above — the padded width the head bank is built with | -| `classifier_control_token_ids` | the control-token ids of the classifier slots; both backends match these in `input_ids` to locate a verdict's read point | +| `probe_num_labels_per_slot` | `len()` of each slot's label token ids (`0` on LoRA slots) | +| `max_probe_labels` | `max()` of the above — the padded width the head bank is built with | +| `probe_control_token_ids` | the control-token ids of the probe slots; both backends match these in `input_ids` to locate a verdict's read point | Note that the label *words* are not stored — only their token ids. Compose resolves and prints the words, but `config.json` keeps ids, so recovering the words means decoding them with the @@ -267,7 +268,7 @@ python -m vllm.entrypoints.openai.api_server \ --port 8000 ``` -Fire a classifier by placing its control token at the very end of the prompt, after the last turn's +Fire a probe by placing its control token at the very end of the prompt, after the last turn's close — one token per slot, `<|safety|>` for the `safety` slot, `<|faithfulness|>` for the other. The verdict is read at the control token, so leave out the generation prompt. The completion is the emitted label word: @@ -286,7 +287,7 @@ curl http://localhost:8000/v1/completions \ {"choices": [{"text": "safe", ...}], ...} ``` -Only the classifier request's final logit row is rewritten; every other (non-classifier) request runs +Only the probe request's final logit row is rewritten; every other (non-probe) request runs untouched with its full vocabulary. If you request logprobs, you will see the label token ids carrying the verdict scores and the rest of the vocabulary pinned at the blanking floor — exactly the two labels for a 2-label slot, three for a 3-label slot. @@ -294,7 +295,7 @@ two labels for a 2-label slot, three for a 3-label slot. ### Firing via the chat template (optional) Instead of placing `<|safety|>` by hand, you can let the composed model's chat template place it for -you by passing `adapter_name` — a classifier slot is activated exactly like any other adapter. The +you by passing `adapter_name` — a probe slot is activated exactly like any other adapter. The template appends the control token after the last turn: ```bash @@ -324,14 +325,14 @@ curl http://localhost:8000/v1/completions \ ``` Only the label ids (`safe`, `unsafe`) will show finite logprobs; all other tokens sit at the blanking -floor. That is the classifier head's per-label verdict surfaced through the standard logprobs +floor. That is the probe head's per-label verdict surfaced through the standard logprobs channel. ## How it works (brief) -- **Routing.** A classifier control token selects its slot like any adapter; its position is routed +- **Routing.** A probe control token selects its slot like any adapter; its position is routed out of the LoRA stream and the head reads the activation there. -- **Padded head bank.** All heads share one tensor padded to `max_classifier_labels` (like the LoRA +- **Padded head bank.** All heads share one tensor padded to `max_probe_labels` (like the LoRA bank pads to `max_lora_rank`), keeping the vLLM head one compile-safe operation. - **Scoring.** Intermediate-layer slots are scored as that layer runs, final-layer slots after the final norm. Scores and slot id travel with the hidden states to the logits step, so this also