diff --git a/README.md b/README.md index 8583f39c..1fdcec23 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,7 @@ starting points: | Vision-language | Gemma 3 | Q4_0 backbone with F16 projector | | Vision-language | Qwen 3.5 | Q4_K_M backbone with BF16 projector | | Vision-language | Qwen 3.6 | UD-IQ2_XXS backbone with BF16 projector | +| Vision-language | Muse Glimmer | Q4_K_XL backbone with Q4_K_M projector (images only) | | Image generation | Z-Image-Turbo | Q4_0 | | Image generation | FLUX.2-klein | Q8_0 | diff --git a/tests/test_multimodal_gguf.py b/tests/test_multimodal_gguf.py index d659d327..01f39188 100644 --- a/tests/test_multimodal_gguf.py +++ b/tests/test_multimodal_gguf.py @@ -32,6 +32,10 @@ class GGUFMMTestConfig(NamedTuple): gguf_model_path: str prompts: list[str] image_names: list[str] + # Per-model, because a model whose image processor emits more patch tokens + # than MAX_MODEL_LEN leaves would have its prompt truncated -- and the + # comparison against HF would then be measuring the truncation. + max_model_len: int = MAX_MODEL_LEN mm_processor_kwargs: dict[str, Any] | None = None @@ -94,6 +98,30 @@ class GGUFMMTestConfig(NamedTuple): image_names=_QWEN35_IMAGE_NAMES, ) +# No ``<|begin_of_text|>``: unlike Gemma 3, this tokenizer's post-processor +# prepends it, so writing it here would produce two. +_MUSE_GLIMMER_PROMPTS = [ + ( + "<|start|>user<|message|><|patch|>" + "What's the content in the center of the image?" + "<|eot|><|start|>assistant" + ), + ("<|start|>user<|message|><|patch|>What is the season?<|eot|><|start|>assistant"), +] + +# The GGUF repo ships neither ``config.json`` nor a tokenizer, so the config has +# to come from the original repo. Naming it as the tokenizer is enough: the +# plugin falls back to the tokenizer path when resolving the GGUF config source. +MUSE_GLIMMER_CONFIG = GGUFMMTestConfig( + original_model="meta-models/Muse-Glimmer-30B", + gguf_model_path="meta-models/Muse-Glimmer-30B-GGUF:Q4_K_XL", + prompts=_MUSE_GLIMMER_PROMPTS, + image_names=_GEMMA3_IMAGE_NAMES, + # The image processor emits up to 4096 patch tokens on its own, so leave + # room for that plus the prompt. + max_model_len=8192, +) + GEMMA3_MODELS_TO_TEST = [ pytest.param(GEMMA3_CONFIG, marks=pytest.mark.slow), pytest.param(GEMMA3_CONFIG_PAN_AND_SCAN, marks=pytest.mark.slow), @@ -112,6 +140,7 @@ def _vllm_generate_greedy_logprobs( max_tokens: int, num_logprobs: int, dtype: str, + max_model_len: int, mm_processor_kwargs: dict[str, Any] | None, ) -> list[tuple[list[int], str, list[dict[int, float] | None]]]: """Run inference via vllm.LLM and return (token_ids, text, logprobs).""" @@ -121,7 +150,7 @@ def _vllm_generate_greedy_logprobs( enforce_eager=True, dtype=dtype, gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - max_model_len=MAX_MODEL_LEN, + max_model_len=max_model_len, mm_processor_kwargs=mm_processor_kwargs, ) try: @@ -260,11 +289,94 @@ def check_logprobs_close( break +# Q4_K noise has no direction, so a shift that does is the signature of a +# conversion the adapter got wrong. Measured across the shipped prompts and +# three image scales, the shared prefix drifts by at most 0.20; a norm that never +# had its offset removed shifts every logprob together and lands far outside +# this. +MAX_LOGPROB_BIAS = 0.5 + + +def check_logprobs_unbiased( + outputs_0_lst: list[tuple[list[int], str, list]], + outputs_1_lst: list[tuple[list[int], str, list]], + name_0: str, + name_1: str, +) -> None: + """Compare the two runs quantitatively wherever that is meaningful. + + Stricter than :func:`check_logprobs_close`, which compares nothing at all + while the tokens agree, tolerates the first divergence with a warning, and + zips the two sequences without checking their lengths -- so an empty run, a + truncated run, and a shared token sequence whose distribution has drifted all + pass it. None of those show up as a load error for Muse Glimmer, whose four + conversions each produce fluent output when they are undone wrongly. + + Comparison stops at the first divergence, and that boundary is the whole + point rather than a shortcut. Greedy decoding conditions each step on what it + already emitted, so up to and including the divergence both runs share a + context and a logprob difference is attributable; past it they are completing + different sentences, and position-wise differences measure that instead. The + gap is not subtle -- the same outputs drift by 0.03 to 0.20 across the shared + prefix and by up to 1.09 once the contexts have parted. + + So the agreeing prefix, which the older comparison skips, is exactly where a + systematic shift is visible, and a bound there is what this adds. + """ + assert len(outputs_0_lst) == len(outputs_1_lst) + + for prompt_idx, (out_0, out_1) in enumerate(zip(outputs_0_lst, outputs_1_lst)): + ids_0, text_0, lps_0 = out_0 + ids_1, text_1, lps_1 = out_1 + context = f"Test {prompt_idx}:\n{name_0}: {text_0!r}\n{name_1}: {text_1!r}" + + assert len(ids_0) == len(ids_1), ( + f"{context}\ngenerated {len(ids_0)} and {len(ids_1)} tokens; a " + "comparison over the shorter of the two would pass on a run that " + "stopped early" + ) + assert ids_0, f"{context}\nboth runs generated nothing" + + diverged = next( + (idx for idx, (a, b) in enumerate(zip(ids_0, ids_1)) if a != b), None + ) + if diverged is None: + shared = len(ids_0) + else: + shared = diverged + 1 + divergence = ( + f"{context}\nfirst differ at token {diverged}: " + f"{name_0} chose {ids_0[diverged]}, {name_1} chose " + f"{ids_1[diverged]}" + ) + # Each side's choice has to at least be a candidate for the other. + # A row permutation left undone picks tokens the reference would + # never rank, which is what this catches. + assert ids_0[diverged] in lps_1[diverged], divergence + assert ids_1[diverged] in lps_0[diverged], divergence + + differences = [ + lps_1[idx][tok] - lps_0[idx][tok] + for idx, tok in enumerate(ids_0[:shared]) + if lps_0 and lps_1 and tok in lps_0[idx] and tok in lps_1[idx] + ] + assert differences, f"{context}\nno position was comparable" + + bias = sum(differences) / len(differences) + assert abs(bias) <= MAX_LOGPROB_BIAS, ( + f"{context}\nmean logprob difference {bias:+.3f} over " + f"{len(differences)} shared positions exceeds {MAX_LOGPROB_BIAS}; " + "quantization noise has no direction, so a drift this one-sided " + "points at a norm offset or a discarded tensor" + ) + + def run_multimodal_gguf_test( model: GGUFMMTestConfig, dtype: str, max_tokens: int, num_logprobs: int, + compare=check_logprobs_close, ) -> None: images = [ImageAsset(name).pil_image for name in model.image_names] size_factors = [0.25, 0.5, 1.0] @@ -286,6 +398,7 @@ def run_multimodal_gguf_test( max_tokens=max_tokens, num_logprobs=num_logprobs, dtype=dtype, + max_model_len=model.max_model_len, mm_processor_kwargs=model.mm_processor_kwargs, ) for prompts, scaled_images in inputs_per_image @@ -304,7 +417,7 @@ def run_multimodal_gguf_test( ] for hf_outputs, gguf_outputs in zip(hf_outputs_per_case, gguf_outputs_per_case): - check_logprobs_close( + compare( outputs_0_lst=hf_outputs, outputs_1_lst=gguf_outputs, name_0="hf", @@ -350,3 +463,126 @@ def test_qwen35_mm_gguf( num_logprobs: int, ) -> None: run_multimodal_gguf_test(model, dtype, max_tokens, num_logprobs) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA required for multimodal GGUF tests.", +) +@pytest.mark.parametrize( + "model", + [pytest.param(MUSE_GLIMMER_CONFIG, marks=pytest.mark.slow)], +) +@pytest.mark.parametrize("dtype", ["bfloat16"]) +@pytest.mark.parametrize("max_tokens", [MAX_TOKENS]) +@pytest.mark.parametrize("num_logprobs", [NUM_LOGPROBS]) +def test_muse_glimmer_mm_gguf( + model: GGUFMMTestConfig, + dtype: str, + max_tokens: int, + num_logprobs: int, +) -> None: + """Images only, compared quantitatively rather than for fluency. + + Video is out of scope here rather than covered: the converter keeps only the + sum of the patch embedding's per-time-step blocks, so it is refused during + input validation instead, and ``test_plugin.py`` is what holds that gate in + place. + + The comparison is :func:`check_logprobs_unbiased` because every conversion + this adapter undoes -- the Q/K row permutation, the norm offset, the + discarded synthetic Q/K norms, the patch embedding split -- loads cleanly + when it is wrong and produces fluent, incorrect text. A comparison that + tolerates the first divergence cannot tell that apart from quantization + noise. + """ + run_multimodal_gguf_test( + model, dtype, max_tokens, num_logprobs, compare=check_logprobs_unbiased + ) + + +def _outputs(tokens: list[int], logprob: float = -0.5, top: float = -0.1): + """One prompt's worth of output, with *tokens* chosen at *logprob*.""" + return [ + ( + tokens, + "".join(chr(97 + tok % 26) for tok in tokens), + [{tok: logprob, -1: top} for tok in tokens], + ) + ] + + +def test_unbiased_check_accepts_an_identical_run(): + reference = _outputs([1, 2, 3, 4]) + + check_logprobs_unbiased(reference, reference, "a", "b") + + +def test_unbiased_check_accepts_noise_without_a_direction(): + """Q4_K noise is what this comparison has to tolerate.""" + left = _outputs([1, 2, 3, 4], logprob=-0.5) + right = [ + ( + left[0][0], + left[0][1], + [ + {tok: lp, -1: -0.1} + for tok, lp in zip(left[0][0], (-0.4, -0.6, -0.45, -0.55)) + ], + ) + ] + + check_logprobs_unbiased(left, right, "a", "b") + + +def test_unbiased_check_rejects_a_run_that_stopped_early(): + """The condition that let a truncated run pass: zip over the shorter side.""" + with pytest.raises(AssertionError, match="stopped early"): + check_logprobs_unbiased(_outputs([1, 2, 3, 4]), _outputs([1, 2]), "a", "b") + + +def test_unbiased_check_rejects_an_empty_run(): + with pytest.raises(AssertionError, match="generated nothing"): + check_logprobs_unbiased(_outputs([]), _outputs([]), "a", "b") + + +def test_unbiased_check_rejects_a_one_sided_logprob_shift(): + """The condition that mattered most: same tokens, drifted distribution. + + A norm whose offset was never removed looks exactly like this -- every + logprob pushed the same way, with the argmax often unchanged. + """ + left = _outputs([1, 2, 3, 4], logprob=-0.5) + right = _outputs([1, 2, 3, 4], logprob=-2.0) + + with pytest.raises(AssertionError, match="one-sided"): + check_logprobs_unbiased(left, right, "a", "b") + + +def test_unbiased_check_rejects_a_choice_the_reference_would_never_rank(): + """What an undone Q/K row permutation produces at the divergence.""" + left = [([1, 2], "a", [{1: -0.1}, {2: -0.5}])] + right = [([1, 9], "b", [{1: -0.1}, {9: -0.5}])] + + with pytest.raises(AssertionError, match="first differ at token"): + check_logprobs_unbiased(left, right, "a", "b") + + +def test_unbiased_check_ignores_drift_once_the_contexts_have_parted(): + """Pins the boundary, which is the part easiest to "strengthen" wrongly. + + Greedy decoding conditions on what it already emitted, so past the divergence + the two runs are completing different sentences. Comparing there reported a + bias of -1.0 on outputs that were both correct descriptions of the image, + which is measuring the divergence rather than the weights. + """ + left = [([1, 2, 3, 4], "a", [{1: -0.1}, {2: -0.5, 9: -0.6}, {3: -0.5}, {4: -0.5}])] + right = [ + ( + [1, 9, 8, 7], + "b", + [{1: -0.1}, {9: -0.5, 2: -0.6}, {8: -0.5, 3: -9.0}, {7: -0.5, 4: -9.0}], + ) + ] + + check_logprobs_unbiased(left, right, "a", "b") diff --git a/tests/test_muse_glimmer_dflash_gguf.py b/tests/test_muse_glimmer_dflash_gguf.py new file mode 100644 index 00000000..268af808 --- /dev/null +++ b/tests/test_muse_glimmer_dflash_gguf.py @@ -0,0 +1,576 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the Muse Glimmer DFlash draft GGUF adapter. + +The draft's adapter is almost entirely negative space: none of the four +conversions the backbone needs applies to it. Its Q/K rows are already in NEOX +order, its norms are stored without the folded offset, and its Q/K norms are +learned rather than synthesized. Applying the backbone's rules here would +rewrite correct weights, and nothing downstream would say so -- the target +verifies every token, so the output stays fluent and correct while the draft's +proposals quietly stop being accepted. Several tests below therefore assert +that a transformation did *not* happen. + +The rest covers the two places where a draft differs structurally from a target +model, both of which produced silent-wrong-weights bugs while this was written: + + · A draft resolves its config and its quantization config separately from the + target, so declarations recorded on the shared objects never reach it. + + · A draft's layers are numbered after the target's, so a module path that + looks right in isolation matches nothing at runtime. + +Everything runs on synthetic names and tensors; none of it needs a checkpoint. +""" + +from types import SimpleNamespace + +import pytest +import torch +from gguf import GGMLQuantizationType, dequantize +from transformers import PretrainedConfig +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, + UnquantizedLinearMethod, +) + +from vllm_gguf_plugin.plugin import _redirect_draft_to_its_config_source +from vllm_gguf_plugin.quantization.config import GGUFConfig +from vllm_gguf_plugin.quantization.linear import GGUFLinearMethod +from vllm_gguf_plugin.quantization.params import ( + _gguf_replicated_weight_loader, + _resolve_gguf_weight_loader, +) +from vllm_gguf_plugin.weights_adapter import ( + get_adapter_architecture, + get_weights_adapter, +) +from vllm_gguf_plugin.weights_adapter.muse_glimmer import ( + MUSE_GLIMMER_DRAFT_ARCHITECTURE, + MuseGlimmerDraftGGUFAdapter, + MuseGlimmerGGUFAdapter, + build_muse_glimmer_draft_name_map, +) + +# One layer's worth of GGUF tensor names, plus the three that live outside the +# blocks. Taken from the shipped dflash checkpoint. +DRAFT_LAYER_TENSORS = [ + "blk.{i}.attn_norm.weight", + "blk.{i}.attn_q.weight", + "blk.{i}.attn_k.weight", + "blk.{i}.attn_v.weight", + "blk.{i}.attn_output.weight", + "blk.{i}.attn_q_norm.weight", + "blk.{i}.attn_k_norm.weight", + "blk.{i}.ffn_norm.weight", + "blk.{i}.ffn_gate.weight", + "blk.{i}.ffn_up.weight", + "blk.{i}.ffn_down.weight", +] +DRAFT_TOP_LEVEL_TENSORS = [ + "fc.weight", + "enc.output_norm.weight", + "output_norm.weight", +] + + +def draft_tensor_names(num_layers: int = 5) -> list[str]: + names = list(DRAFT_TOP_LEVEL_TENSORS) + for layer in range(num_layers): + names += [name.format(i=layer) for name in DRAFT_LAYER_TENSORS] + return names + + +def draft_config(model_type: str = "muse_glimmer_assistant") -> PretrainedConfig: + return PretrainedConfig(model_type=model_type) + + +def eagle_wrapped(inner: PretrainedConfig) -> PretrainedConfig: + """Mimic how EAGLEConfig presents a draft config to the loader. + + It reports ``model_type == "eagle"`` and keeps the real config on ``.model``, + and it deliberately does not copy the inner ``model_type`` up. + """ + return PretrainedConfig(model_type="eagle", model=inner) + + +# -------------------------------------------------------------------------- +# Adapter selection +# -------------------------------------------------------------------------- +def test_the_draft_adapter_claims_the_assistant_config(): + assert isinstance(get_weights_adapter(draft_config()), MuseGlimmerDraftGGUFAdapter) + + +def test_the_draft_adapter_is_still_found_through_the_eagle_wrapper(): + """The wrapper goes on after the config parser has run. + + An adapter that only checks the bare model type matches while the + architecture is being chosen and stops matching by the time weights are + mapped, at which point the fallback adapter takes over and fails on an + architecture it has never heard of. + """ + wrapped = eagle_wrapped(draft_config()) + assert isinstance(get_weights_adapter(wrapped), MuseGlimmerDraftGGUFAdapter) + + +def test_the_draft_adapter_declares_the_bare_architecture(): + """EAGLEConfig rewrites this to DFlash{arch}; it must not be pre-rewritten.""" + assert get_adapter_architecture(draft_config()) == MUSE_GLIMMER_DRAFT_ARCHITECTURE + assert not MUSE_GLIMMER_DRAFT_ARCHITECTURE.startswith("DFlash") + + +@pytest.mark.parametrize("model_type", ["muse_glimmer", "muse_glimmer_text"]) +def test_the_draft_adapter_leaves_the_backbone_alone(model_type): + config = PretrainedConfig(model_type=model_type) + assert not MuseGlimmerDraftGGUFAdapter.matches(config) + assert isinstance(get_weights_adapter(config), MuseGlimmerGGUFAdapter) + + +def test_the_backbone_adapter_does_not_claim_the_draft(): + assert not MuseGlimmerGGUFAdapter.matches(draft_config()) + + +def test_an_unrelated_eagle_draft_is_not_claimed(): + """The wrapper is generic, so unwrapping must not widen what matches.""" + wrapped = eagle_wrapped(PretrainedConfig(model_type="llama")) + assert not MuseGlimmerDraftGGUFAdapter.matches(wrapped) + + +# -------------------------------------------------------------------------- +# Name mapping +# -------------------------------------------------------------------------- +def test_the_name_map_is_a_bijection_over_the_checkpoint(): + names = draft_tensor_names() + name_map = build_muse_glimmer_draft_name_map(names) + + assert set(name_map) == set(names), "some GGUF tensor went unmapped" + assert len(set(name_map.values())) == len(names), "two tensors share a name" + + +def test_the_draft_discards_nothing(): + """Unlike the backbone, which drops 104 synthesized Q/K norms. + + Reusing that rule here would remove real learned weights. + """ + names = draft_tensor_names() + name_map = build_muse_glimmer_draft_name_map(names) + + assert len(name_map) == len(names) + for layer in range(5): + for kind in ("q", "k"): + assert f"blk.{layer}.attn_{kind}_norm.weight" in name_map + + +@pytest.mark.parametrize( + "gguf_name,hf_name", + [ + ("fc.weight", "encoder.fc.weight"), + ("enc.output_norm.weight", "encoder.output_norm_enc.weight"), + ("output_norm.weight", "norm.weight"), + ], +) +def test_the_tensors_outside_the_blocks_are_renamed(gguf_name, hf_name): + """llama.cpp abbreviates these three; nothing about them is derivable.""" + name_map = build_muse_glimmer_draft_name_map(draft_tensor_names()) + assert name_map[gguf_name] == hf_name + + +@pytest.mark.parametrize( + "suffix,hf_suffix", + [ + ("attn_q", "self_attn.q_proj"), + ("attn_k", "self_attn.k_proj"), + ("attn_v", "self_attn.v_proj"), + ("attn_output", "self_attn.o_proj"), + ("attn_q_norm", "self_attn.q_norm"), + ("attn_k_norm", "self_attn.k_norm"), + ("attn_norm", "input_layernorm"), + ("ffn_norm", "post_attention_layernorm"), + ("ffn_gate", "mlp.gate_proj"), + ("ffn_up", "mlp.up_proj"), + ("ffn_down", "mlp.down_proj"), + ], +) +def test_the_within_layer_renames(suffix, hf_suffix): + name_map = build_muse_glimmer_draft_name_map(draft_tensor_names()) + assert name_map[f"blk.3.{suffix}.weight"] == f"layers.3.{hf_suffix}.weight" + + +def test_an_unrecognized_tensor_is_skipped_rather_than_guessed(caplog): + name_map = build_muse_glimmer_draft_name_map( + [*draft_tensor_names(1), "blk.0.something_new.weight"] + ) + assert "blk.0.something_new.weight" not in name_map + assert "something_new" in caplog.text + + +# -------------------------------------------------------------------------- +# What the draft must *not* do to its weights +# -------------------------------------------------------------------------- +def build_adapter(num_layers: int = 2) -> MuseGlimmerDraftGGUFAdapter: + """Create an adapter in the order the loader calls it.""" + adapter = MuseGlimmerDraftGGUFAdapter() + adapter.build_name_map( + SimpleNamespace(all_files=(), backbone=(), primary_backbone=None), + SimpleNamespace(hf_config=draft_config(), dtype=torch.bfloat16), + ) + return adapter + + +@pytest.fixture +def draft_adapter(monkeypatch): + monkeypatch.setattr( + "vllm_gguf_plugin.weights_adapter.muse_glimmer.get_gguf_tensor_names", + lambda _files: draft_tensor_names(2), + ) + return build_adapter() + + +def transformed(adapter, weights): + model_config = SimpleNamespace(hf_config=draft_config(), dtype=torch.bfloat16) + return dict(adapter.transform_weights(iter(weights), model_config)) + + +# A Q4_K super-block spends 144 bytes on 256 weights: the fp16 pair `d`/`dmin`, +# then 6-bit scales and 4-bit quants, both integer fields. +Q4_K_WEIGHTS_PER_BLOCK = 256 +Q4_K_BLOCK_BYTES = 144 + + +def packed_q4_k(rows: int, cols: int, *, seed: int) -> torch.Tensor: + """Random Q4_K bytes whose super-block scales are finite. + + The opening pair is the only part read as floating point, and uniform + random bytes give one of them a NaN exponent for about one super-block in + twenty. That would be a test failing on its own data: a NaN compares + unequal to itself, so the dequantized rows would not match a reference + dequantization of the very same bytes. + """ + blocks = cols // Q4_K_WEIGHTS_PER_BLOCK + generator = torch.Generator().manual_seed(seed) + packed = torch.randint( + 0, + 256, + (rows, blocks * Q4_K_BLOCK_BYTES), + dtype=torch.uint8, + generator=generator, + ) + scales = torch.tensor([1.0, 0.5], dtype=torch.float16).view(torch.uint8) + for block in range(blocks): + start = block * Q4_K_BLOCK_BYTES + packed[:, start : start + scales.numel()] = scales + return packed + + +def test_the_norms_pass_through_untouched(draft_adapter): + """The backbone subtracts one here. The draft must not. + + Its norms are plain RMSNorm weights, so an offset that the backbone's + checkpoint folds in was never folded in to begin with. + """ + norm = torch.tensor([1.5, 0.5, 2.0], dtype=torch.bfloat16) + out = transformed(draft_adapter, [("layers.0.input_layernorm.weight", norm)]) + + torch.testing.assert_close(out["layers.0.input_layernorm.weight"], norm) + + +def test_the_final_norm_passes_through_untouched(draft_adapter): + norm = torch.tensor([1.5, 0.5, 2.0], dtype=torch.bfloat16) + out = transformed(draft_adapter, [("norm.weight", norm)]) + + torch.testing.assert_close(out["norm.weight"], norm) + + +def test_the_qk_norms_survive(draft_adapter): + """They are learned weights here, not the synthesized ones the backbone drops.""" + weight = torch.tensor([1.25, 0.75], dtype=torch.bfloat16) + out = transformed(draft_adapter, [("layers.1.self_attn.q_norm.weight", weight)]) + + torch.testing.assert_close(out["layers.1.self_attn.q_norm.weight"], weight) + + +def test_the_output_projection_keeps_its_packed_bytes(draft_adapter): + """Only Q/K/V are unpacked; everything else stays quantized. + + Unpacking more than necessary is not a correctness bug, which is exactly why + it needs a test -- it would just quietly cost memory. + """ + packed = torch.arange(16, dtype=torch.uint8) + out = transformed( + draft_adapter, + [ + ("layers.0.self_attn.o_proj.qweight_type", torch.tensor(14)), + ("layers.0.self_attn.o_proj.qweight", packed), + ], + ) + + assert "layers.0.self_attn.o_proj.weight" not in out + torch.testing.assert_close(out["layers.0.self_attn.o_proj.qweight"], packed) + + +@pytest.mark.parametrize("kind", ["q", "k", "v"]) +def test_the_qkv_projections_are_unpacked(kind, draft_adapter): + """The head reads qkv_proj.weight to build its fused KV buffer.""" + rows, cols = 4, 256 + block = GGMLQuantizationType.Q4_K + packed = packed_q4_k(rows, cols, seed=0) + name = f"layers.0.self_attn.{kind}_proj" + out = transformed( + draft_adapter, + [ + (f"{name}.qweight_type", torch.tensor(int(block))), + (f"{name}.qweight", packed), + ], + ) + + assert f"{name}.qweight" not in out, "the packed form must not also be yielded" + assert out[f"{name}.weight"].shape == (rows, cols) + assert out[f"{name}.weight"].dtype == torch.bfloat16 + + +@pytest.mark.parametrize("kind", ["q", "k"]) +def test_the_qk_rows_keep_the_order_they_were_stored_in(kind, draft_adapter): + """The backbone reorders these rows. The draft's are already NEOX. + + Unpacking is the only thing allowed to happen to them, so the result has to + equal a plain dequantization of the same bytes. A permutation added here + would not fail loudly: the target verifies every token, so the output would + stay fluent while the draft's proposals stopped being accepted. + """ + rows, cols = 8, 256 + block = GGMLQuantizationType.Q4_K + packed = packed_q4_k(rows, cols, seed=1) + name = f"layers.0.self_attn.{kind}_proj" + out = transformed( + draft_adapter, + [ + (f"{name}.qweight_type", torch.tensor(int(block))), + (f"{name}.qweight", packed), + ], + ) + + reference = torch.from_numpy( + dequantize(packed.numpy(), block).reshape(rows, cols) + ).to(torch.bfloat16) + assert reference.isfinite().all(), ( + "a NaN here would compare unequal to itself and fail the assertion below " + "regardless of what the adapter did" + ) + torch.testing.assert_close(out[f"{name}.weight"], reference) + + +# -------------------------------------------------------------------------- +# Reaching a draft's own quantization config +# -------------------------------------------------------------------------- +@pytest.fixture +def single_rank(monkeypatch): + """Let linear layers be built without a distributed group. + + Only the layer's type matters to the code under test, but constructing one + still asks for the tensor-parallel rank. + """ + import vllm.model_executor.layers.linear as linear_module + import vllm.model_executor.parameter as parameter_module + + for module in (linear_module, parameter_module): + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1) + + +def dense_config() -> GGUFConfig: + config = GGUFConfig( + dense_module_suffixes=list(MuseGlimmerDraftGGUFAdapter.dense_module_suffixes) + ) + config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + return config + + +def test_a_declaration_survives_being_rebuilt_from_the_config_dict(): + """A draft's layers are built against a config rebuilt from this dict. + + Anything the loader records on the shared object reaches the target model + and never the draft, so dropping these keys would leave a draft with an + empty declaration and no sign of why. + """ + rebuilt = GGUFConfig.from_config( + { + "quant_method": "gguf", + "unquantized_modules": ["embed_tokens"], + "dense_module_suffixes": ["self_attn.qkv_proj"], + } + ) + + assert rebuilt.unquantized_modules == ["embed_tokens"] + assert rebuilt.dense_module_suffixes == ["self_attn.qkv_proj"] + + +def test_a_config_dict_without_the_keys_still_builds(): + """Target models supply neither key.""" + rebuilt = GGUFConfig.from_config({"quant_method": "gguf"}) + + assert rebuilt.unquantized_modules == [] + assert rebuilt.dense_module_suffixes == [] + + +def test_a_suffix_declaration_survives_the_layer_renumbering(single_rank): + """vLLM numbers a draft's layers after the target's. + + A five-layer draft behind a sixty-two-layer target is asked about + ``model.layers.62..66``. Declaring the layers by name means predicting that + offset; matching on the suffix does not. + """ + config = dense_config() + layer = QKVParallelLinear( + hidden_size=64, + head_size=16, + total_num_heads=4, + quant_config=None, + disable_tp=True, + ) + + for index in (0, 62, 66, 999): + method = config.get_quant_method( + layer, prefix=f"model.layers.{index}.self_attn.qkv_proj" + ) + assert isinstance(method, UnquantizedLinearMethod), ( + f"layer {index} was left quantized" + ) + + +def test_the_suffix_declaration_does_not_leak_to_other_layers(single_rank): + """Only the fused attention projection is exempt.""" + config = dense_config() + layer = RowParallelLinear( + input_size=64, output_size=64, quant_config=None, disable_tp=True + ) + + for suffix in ("self_attn.o_proj", "mlp.down_proj", "mlp.gate_up_proj", "fc"): + method = config.get_quant_method(layer, prefix=f"model.layers.3.{suffix}") + assert isinstance(method, GGUFLinearMethod), f"{suffix} should stay packed" + + +def test_a_config_without_suffixes_quantizes_everything(single_rank): + """Negative control: the exemption comes from the declaration, not the path.""" + config = GGUFConfig() + config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + layer = QKVParallelLinear( + hidden_size=64, + head_size=16, + total_num_heads=4, + quant_config=None, + disable_tp=True, + ) + + method = config.get_quant_method(layer, prefix="model.layers.62.self_attn.qkv_proj") + assert isinstance(method, GGUFLinearMethod) + + +# -------------------------------------------------------------------------- +# The one linear layer GGUF had never had to load +# -------------------------------------------------------------------------- +def test_an_unsharded_linear_gets_a_loader_that_can_size_itself(single_rank): + """ReplicatedLinear has no v2 loader, and its v1 loader asserts the shape. + + GGUF parameters start empty and take their shape from the packed bytes, so + the assertion fires on the first tensor. The draft's ``fc`` is the first + ReplicatedLinear any GGUF model has needed. + """ + layer = ReplicatedLinear( + input_size=32, output_size=8, quant_config=None, disable_tp=True + ) + + assert not hasattr(layer, "weight_loader_v2") + resolved = _resolve_gguf_weight_loader(layer, layer.weight_loader) + assert resolved is _gguf_replicated_weight_loader + + +def test_a_sharded_linear_still_uses_the_v2_loader(single_rank): + layer = RowParallelLinear( + input_size=32, output_size=8, quant_config=None, disable_tp=True + ) + + resolved = _resolve_gguf_weight_loader(layer, layer.weight_loader) + assert resolved == layer.weight_loader_v2 + + +# -------------------------------------------------------------------------- +# Pointing a separate-file draft at its own config +# -------------------------------------------------------------------------- +def engine_args(speculative_config): + return SimpleNamespace(speculative_config=speculative_config) + + +def test_a_draft_is_redirected_to_the_config_directory_it_was_given(tmp_path): + config_dir = tmp_path / "assistant" + config_dir.mkdir() + (config_dir / "config.json").write_text("{}") + draft = str(tmp_path / "draft.gguf") + + args = engine_args({"model": draft, "hf_config_path": str(config_dir)}) + weights = _redirect_draft_to_its_config_source(args) + + assert weights == draft, "the file has to come back as the weights source" + assert args.speculative_config["model"] == str(config_dir) + assert args.speculative_config["quantization"] == "gguf" + assert "hf_config_path" not in args.speculative_config, ( + "SpeculativeConfig rejects fields it does not declare" + ) + + +def test_redirecting_twice_still_reports_the_weights(tmp_path): + """The rewrite edits the caller's dict, so a second pass sees no GGUF path. + + Losing the weights path there is not an error: the draft keeps the config + directory as its weights source and loads whatever checkpoint is sitting in + it, which for this layout is the unquantized one. + """ + config_dir = tmp_path / "assistant" + config_dir.mkdir() + (config_dir / "config.json").write_text("{}") + draft = str(tmp_path / "draft.gguf") + args = engine_args({"model": draft, "hf_config_path": str(config_dir)}) + + first = _redirect_draft_to_its_config_source(args) + second = _redirect_draft_to_its_config_source(args) + + assert first == second == draft + + +def test_a_draft_without_a_config_says_what_to_pass(tmp_path): + """The directory a draft sits in belongs to the model it drafts for.""" + draft = tmp_path / "draft.gguf" + draft.write_bytes(b"") + args = engine_args({"model": str(draft)}) + + with pytest.raises(ValueError, match="hf_config_path"): + _redirect_draft_to_its_config_source(args) + + +def test_an_unquantized_draft_is_left_alone(): + args = engine_args({"model": "/models/some-draft", "num_speculative_tokens": 3}) + + assert _redirect_draft_to_its_config_source(args) is None + assert args.speculative_config["model"] == "/models/some-draft" + assert "quantization" not in args.speculative_config + + +def test_no_speculative_config_is_left_alone(): + assert _redirect_draft_to_its_config_source(engine_args(None)) is None + + +def test_an_explicit_quantization_choice_is_respected(tmp_path): + config_dir = tmp_path / "assistant" + config_dir.mkdir() + (config_dir / "config.json").write_text("{}") + args = engine_args( + { + "model": str(tmp_path / "draft.gguf"), + "hf_config_path": str(config_dir), + "quantization": "awq", + } + ) + + _redirect_draft_to_its_config_source(args) + + assert args.speculative_config["quantization"] == "awq" diff --git a/tests/test_muse_glimmer_gguf.py b/tests/test_muse_glimmer_gguf.py new file mode 100644 index 00000000..c12ad8c7 --- /dev/null +++ b/tests/test_muse_glimmer_gguf.py @@ -0,0 +1,536 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the Muse Glimmer GGUF adapter. + +Every conversion this adapter undoes fails quietly when it is wrong: the weights +load without complaint and the model generates fluent, incorrect text. So these +tests pin the conversions themselves rather than that loading succeeds. They run +on synthetic tensors and synthetic tensor names, so none of them needs the real +30B checkpoint. + +Muse Glimmer GGUF checkpoints store Q/K with llama.cpp's interleaved rotary +layout while vLLM's implementation hardcodes NEOX, so the adapter permutes rows +on the way in. The first group of tests pins the three properties that make the +permutation safe to apply to packed quantized bytes. + +Quantized coverage uses the small synthetic sample tensors from +``Isotr0py/test-gguf-sample`` rather than a real checkpoint, since the property +under test is generic to any K-quant tensor. + +The dequantization reference is the Triton kernel, deliberately: it computes in +fp32 and matches ``gguf.dequantize`` bit-for-bit, whereas the native CUDA/HIP +kernel computes in fp16 (see ``csrc/gguf/dequantize.cuh``) and would turn these +exact comparisons into tolerance comparisons. +""" + +from types import SimpleNamespace + +import pytest +import torch +from gguf import GGML_QUANT_SIZES, GGMLQuantizationType +from transformers import PretrainedConfig + +from vllm_gguf_plugin.gguf_files import GGUFModelFiles +from vllm_gguf_plugin.triton.dequantize.interface import ggml_dequantize_triton +from vllm_gguf_plugin.weights_adapter import muse_glimmer as adapter_module +from vllm_gguf_plugin.weights_adapter.muse_glimmer import ( + MUSE_GLIMMER_ARCHITECTURES, + MuseGlimmerGGUFAdapter, + has_vision, + interleaved_to_neox_row_index, + neox_to_interleaved_row_index, + reconstruct_patch_embedding, + undo_rope_interleave, +) + +from .utils import get_gguf_sample_tensors + +# 96 and 128 are the vision and text head_dim of the shipped Muse Glimmer +# checkpoints; the rest are there to catch off-by-one indexing. 4 is excluded on +# purpose: the permutation is accidentally self-inverse there, which hides bugs. +HEAD_DIMS = [8, 16, 64, 96, 128] +K_QUANT_TYPES = [ + GGMLQuantizationType.Q4_K, + GGMLQuantizationType.Q5_K, + GGMLQuantizationType.Q6_K, +] +HIDDEN_SIZES = [256, 1024] + + +@pytest.mark.parametrize("head_dim", HEAD_DIMS) +@pytest.mark.parametrize("num_heads", [1, 2, 32]) +def test_forward_and_inverse_index_compose_to_identity(num_heads, head_dim): + forward = neox_to_interleaved_row_index(num_heads, head_dim) + inverse = interleaved_to_neox_row_index(num_heads, head_dim) + identity = torch.arange(num_heads * head_dim) + + torch.testing.assert_close(forward[inverse], identity) + torch.testing.assert_close(inverse[forward], identity) + + +@pytest.mark.parametrize("head_dim", [8, 16, 64, 128]) +def test_inverse_index_is_not_self_inverse(head_dim): + """Guards the most likely way to get this wrong. + + Applying the permutation twice looks like an inverse for head_dim == 4 but + diverges for every larger size, so a helper that calls the forward function + twice passes toy tests and corrupts real checkpoints. + """ + inverse = interleaved_to_neox_row_index(2, head_dim) + identity = torch.arange(2 * head_dim) + + assert not torch.equal(inverse[inverse], identity) + + +def test_head_dim_four_is_the_misleading_case(): + """Documents why head_dim == 4 must not be used as the only test size.""" + inverse = interleaved_to_neox_row_index(1, 4) + + assert torch.equal(inverse[inverse], torch.arange(4)) + + +@pytest.mark.parametrize("head_dim", HEAD_DIMS) +def test_undo_rope_interleave_inverts_the_conversion(head_dim): + """End-to-end on float rows: forward re-layout then undo is the identity.""" + num_heads = 4 + original = torch.randn(num_heads * head_dim, 7) + + converted = original.index_select( + 0, neox_to_interleaved_row_index(num_heads, head_dim) + ) + assert not torch.equal(converted, original) + + torch.testing.assert_close( + undo_rope_interleave(converted, num_heads, head_dim), original + ) + + +def test_undo_rope_interleave_rejects_row_count_mismatch(): + with pytest.raises(ValueError, match="expected 256 rows"): + undo_rope_interleave(torch.zeros(255, 4), num_heads=2, head_dim=128) + + +def test_odd_head_dim_is_rejected(): + with pytest.raises(ValueError, match="head_dim must be even"): + interleaved_to_neox_row_index(1, 7) + + +@pytest.mark.parametrize("quant_type", K_QUANT_TYPES) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +def test_qweight_rows_are_self_contained_byte_runs(quant_type, hidden_size): + """Super-blocks split along the input dim, so a row is whole blocks. + + This is the precondition that lets the adapter permute packed bytes at all. + """ + block_size, type_size = GGML_QUANT_SIZES[quant_type] + + for tensor in get_gguf_sample_tensors(hidden_size, quant_type): + qweight = torch.tensor(tensor.data) + assert qweight.ndim == 2 + row_bytes = qweight.shape[1] + + assert row_bytes % type_size == 0, ( + f"{tensor.name}: row of {row_bytes} bytes is not a whole number of " + f"{type_size}-byte blocks" + ) + assert (row_bytes // type_size) * block_size == hidden_size + + +@pytest.mark.parametrize("quant_type", K_QUANT_TYPES) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +def test_qweight_byte_permutation_round_trips(quant_type, hidden_size): + """Permuting packed bytes forward then back recovers the original bytes.""" + head_dim = 128 + + for tensor in get_gguf_sample_tensors(hidden_size, quant_type): + qweight = torch.tensor(tensor.data) + num_rows = qweight.shape[0] + if num_rows % head_dim: + continue + num_heads = num_rows // head_dim + + forward = qweight.index_select( + 0, neox_to_interleaved_row_index(num_heads, head_dim) + ) + restored = undo_rope_interleave(forward, num_heads, head_dim) + + assert torch.equal(restored, qweight), f"{tensor.name}: bytes changed" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("quant_type", K_QUANT_TYPES) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +def test_permute_then_dequantize_equals_dequantize_then_permute( + quant_type, hidden_size +): + """The load-bearing equivalence: reordering bytes is safe. + + If this holds, the adapter can permute rows in the quantized domain at zero + memory cost and zero precision loss, instead of dequantizing Q/K (which + would drag the whole fused QKV projection to bf16). + """ + head_dim = 128 + block_size, type_size = GGML_QUANT_SIZES[quant_type] + checked = 0 + + for tensor in get_gguf_sample_tensors(hidden_size, quant_type): + qweight = torch.tensor(tensor.data).cuda() + num_rows, row_bytes = qweight.shape + if num_rows % head_dim: + continue + num_heads = num_rows // head_dim + shape = (num_rows, row_bytes // type_size * block_size) + + straight = ggml_dequantize_triton( + qweight, int(quant_type), *shape, torch.float32 + ) + permuted_first = ggml_dequantize_triton( + undo_rope_interleave(qweight, num_heads, head_dim), + int(quant_type), + *shape, + torch.float32, + ) + dequantized_first = undo_rope_interleave(straight, num_heads, head_dim) + + # Without this the equality below would also hold for a no-op permutation. + assert not torch.equal(permuted_first, straight), ( + f"{tensor.name}: permutation had no effect, equality is vacuous" + ) + assert torch.equal(permuted_first, dequantized_first), ( + f"{tensor.name}: permuting bytes and permuting floats disagree" + ) + checked += 1 + + assert checked, "no sample tensor had a row count divisible by head_dim" + + +# -------------------------------------------------------------------------- +# Which weights the adapter loads at all +# -------------------------------------------------------------------------- +VISION_LAYERS = (0, 1) +# The GGUF side of the naming, which belongs to the file format rather than to +# this plugin, so it is spelled out here instead of read back off the mapper. +_VISION_LEAVES = ("attn_q", "attn_k", "attn_v", "attn_out") + + +def _vision_config(**kwargs) -> PretrainedConfig: + config = PretrainedConfig() + config.vision_config = PretrainedConfig() + for key, value in kwargs.items(): + setattr(config, key, value) + return config + + +def test_has_vision_matches_the_rule_vllm_applies(): + """The two have to stay the same rule, so compare against vLLM's own. + + The adapter decides which weights to produce and vLLM decides which modules + to build. Whenever the two disagree, one side has weights the other has + nowhere to put -- so this pins them together rather than trusting that a + copied predicate stays a copy. + """ + from vllm.model_executor.models.muse_glimmer import _muse_glimmer_has_vision + + configs = [ + _vision_config(), + _vision_config(has_vision=True), + _vision_config(has_vision=False), + PretrainedConfig(), + ] + + for config in configs: + assert has_vision(config) == _muse_glimmer_has_vision(config) + + # Spelled out as well, so a change to both at once still has to be deliberate. + assert [has_vision(config) for config in configs] == [True, True, False, False] + + +def test_architecture_names_the_class_vllm_will_build(): + """Crossing the two model types is another failure that does not raise. + + Neither type can be looked up in the mapping the config parser consults, so + the adapter has to name the class itself. Pointing the multimodal type at + the causal-LM class loads and generates perfectly well -- as a model with no + vision tower at all -- so this checks both against their real sources rather + than against a copy of the same dict. + """ + from transformers.models.auto.modeling_auto import ( + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, + ) + from vllm.model_executor.models.registry import ModelRegistry + + assert ( + MUSE_GLIMMER_ARCHITECTURES["muse_glimmer"] + == (MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES["muse_glimmer"]) + ) + assert ( + MUSE_GLIMMER_ARCHITECTURES["muse_glimmer_text"] + != (MUSE_GLIMMER_ARCHITECTURES["muse_glimmer"]) + ) + + supported = ModelRegistry.get_supported_archs() + for model_type, architecture in MUSE_GLIMMER_ARCHITECTURES.items(): + config = SimpleNamespace(model_type=model_type) + assert MuseGlimmerGGUFAdapter.matches(config) + assert MuseGlimmerGGUFAdapter.architecture(config) == architecture + assert architecture in supported + + +def _build(monkeypatch, config, *, projector: str | None, gguf_names=()): + """Drive the adapter through the loader's sequence, without real files.""" + monkeypatch.setattr( + adapter_module, + "maybe_patch_hf_config_from_gguf", + lambda _path, cfg, mmproj_path=None: cfg, + ) + monkeypatch.setattr( + adapter_module, "get_gguf_tensor_names", lambda _files: list(gguf_names) + ) + + files = GGUFModelFiles(backbone=("backbone.gguf",), mm_proj=projector) + adapter = MuseGlimmerGGUFAdapter() + patched = adapter.patch_hf_config(files, config) + name_map = adapter.build_name_map(files, SimpleNamespace(hf_config=patched)) + return adapter, name_map + + +_PROJECTOR_NAMES = ("v.blk.0.attn_q.weight", "mm.2.weight") + + +def test_projector_weights_are_mapped_when_vision_is_on(monkeypatch): + _, name_map = _build( + monkeypatch, + _vision_config(), + projector="mmproj.gguf", + gguf_names=_PROJECTOR_NAMES, + ) + + assert sorted(name_map) == sorted(_PROJECTOR_NAMES) + + +def test_missing_projector_fails_instead_of_loading_half_a_model(monkeypatch): + """A silent partial load is the failure worth preventing here. + + Requesting one quantization of a repo does not bring the projector along, + since it is quantized separately. The config still describes a vision tower, + so vLLM builds one and simply never receives its weights -- which text + prompts do not reveal, and image prompts reveal only as bad output. + """ + with pytest.raises(RuntimeError, match="mmproj"): + _build(monkeypatch, _vision_config(), projector=None) + + +@pytest.mark.parametrize( + "config", + [ + pytest.param(_vision_config(has_vision=False), id="vision_turned_off"), + pytest.param(PretrainedConfig(), id="text_only_config"), + ], +) +def test_projector_weights_are_left_out_when_vision_is_off(monkeypatch, config): + """Text-only runs must not pick the projector up off the directory. + + The loader resolves the projector from the directory, and downloads one when + the config carries a ``vision_config``, so it arrives whether or not this run + wants it. vLLM builds no vision tower for either config here, so leaving + those tensors out of the map is what keeps 809 weights from arriving with + nowhere to go. + """ + _, name_map = _build( + monkeypatch, config, projector="mmproj.gguf", gguf_names=_PROJECTOR_NAMES + ) + + assert name_map == {} + + +# -------------------------------------------------------------------------- +# Name mapping +# -------------------------------------------------------------------------- +def test_synthesized_qk_norms_are_discarded(monkeypatch, caplog): + """These are not learned parameters and loading them applies a factor twice. + + The converter synthesizes them as constants from the config's scale factor, + which vLLM already applies itself. They are also checked not to be reported + as unmapped: that warning is for names the mapping failed to cover, and + listing a deliberate omission there would send the next reader after a + mapping bug that does not exist. + """ + names = [f"blk.0.attn_{proj}_norm.weight" for proj in ("q", "k")] + + _, name_map = _build( + monkeypatch, _vision_config(), projector="mmproj.gguf", gguf_names=names + ) + + assert name_map == {} + assert "No HF name" not in caplog.text + + +def test_text_and_vision_projections_are_told_apart(monkeypatch): + """Both towers call the projection ``attn_q`` in GGUF; the targets differ. + + The text module ``self_attn.q_proj`` ends with the vision module's + ``attn.q_proj``, so a substring rule alone would rewrite vision names into + text ones. + """ + _, name_map = _build( + monkeypatch, + _vision_config(), + projector="mmproj.gguf", + gguf_names=["blk.0.attn_q.weight", "v.blk.0.attn_q.weight"], + ) + + assert name_map == { + "blk.0.attn_q.weight": ( + "model.language_model.layers.0.self_attn.q_proj.weight" + ), + "v.blk.0.attn_q.weight": "model.vision_tower.layers.0.attn.q_proj.weight", + } + + +# -------------------------------------------------------------------------- +# Declaring the dequantized modules unquantized +# -------------------------------------------------------------------------- +def test_declaration_cannot_be_read_before_the_name_map_is_built(): + """It is derived from the name map, so reading it early would under-declare. + + The loader builds the map first today. Were the declaration assigned during + that call rather than derived from it, a reordering on the loader's side + would quietly return an empty list -- and an under-declared fused layer + allocates packed buffers that no incoming tensor fills. + """ + adapter = MuseGlimmerGGUFAdapter() + + with pytest.raises(RuntimeError, match="before build_name_map"): + _ = adapter.extra_unquantized_modules + + +def _declaration_as_vllm_sees_it(monkeypatch) -> tuple[list[str], dict[str, list[str]]]: + """What the adapter declared, in vLLM's namespace. + + Taken from the adapter rather than assembled here, so that dropping either + half of the declaration shows up as a failure instead of leaving the test + agreeing with itself. + + Three steps at runtime: the loader records HF names, vLLM rewrites them with + ``apply_vllm_mapper``, and the prefix ``get_quant_method`` receives is a vLLM + name too. Nearly every name changes on the way through -- the text tower + loses ``model.language_model.`` and ``model.vision_tower.`` becomes + ``vision_encoder.`` -- so a declaration that only agrees with itself in HF + names says nothing about what happens at load time. + + The loader's own half of the declaration -- the modules it derives from the + tensors a GGUF stores unquantized -- is left out on purpose, which is what + passing only synthetic names achieves. Real checkpoints carry attention + biases among those, whose full paths cover the fused ``qkv_proj`` for free; + relying on that is what kept ``o_proj`` looking covered while it was not. + """ + from vllm.model_executor.models.muse_glimmer import MuseGlimmerForCausalLM + + from vllm_gguf_plugin.quantization.config import GGUFConfig + + gguf_names = [ + f"v.blk.{layer}.{leaf}.weight" + for layer in VISION_LAYERS + for leaf in _VISION_LEAVES + ] + ["token_embd.weight", "mm.2.weight"] + + adapter, _ = _build( + monkeypatch, + _vision_config(), + projector="mmproj.gguf", + gguf_names=gguf_names, + ) + + config = GGUFConfig(unquantized_modules=list(adapter.extra_unquantized_modules)) + config.apply_vllm_mapper( + MuseGlimmerForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper() + ) + return config.unquantized_modules, MuseGlimmerForCausalLM.packed_modules_mapping + + +@pytest.mark.parametrize("layer", VISION_LAYERS) +@pytest.mark.parametrize("projection", ["qkv_proj", "o_proj"]) +def test_dequantized_vision_layers_are_declared_unquantized( + monkeypatch, layer, projection +): + """Whatever is dequantized has to be judged unquantized -- one decision. + + The adapter hands these modules plain float ``weight``. If + ``get_quant_method`` thinks they are quantized it allocates ``qweight`` + buffers that no incoming tensor fills. + + Nothing reaches that branch today, because vLLM builds these layers without + forwarding ``quant_config`` -- which is what forced the dequantization in the + first place. So this guards the day vLLM forwards it, when a missing + declaration turns into a hard error. + + The two projections are covered by different halves of the declaration, which + is why both are checked: ``qkv_proj`` is fused, and the fused branch of + ``is_layer_skipped_gguf`` asks whether a declared name *contains* the shard's + full path, so only a full-depth name matches it -- while ``o_proj`` is + renamed on vLLM's way in (from ``attn.proj``) and is reached only by the + coarse prefix. + """ + from vllm_gguf_plugin.quantization.utils import is_layer_skipped_gguf + + declared, fused_mapping = _declaration_as_vllm_sees_it(monkeypatch) + prefix = f"vision_encoder.transformer.{layer}.attn.{projection}" + + assert is_layer_skipped_gguf(prefix, declared, fused_mapping), ( + f"{prefix} is handed a dequantized float weight but is not declared " + "unquantized; once vLLM forwards quant_config it would allocate qweight " + "buffers nothing fills" + ) + + +def test_embedding_is_declared_unquantized(monkeypatch): + from vllm_gguf_plugin.quantization.utils import is_layer_skipped_gguf + + declared, fused_mapping = _declaration_as_vllm_sees_it(monkeypatch) + + assert is_layer_skipped_gguf("model.embed_tokens", declared, fused_mapping) + + +@pytest.mark.parametrize( + "prefix", + [ + "model.layers.0.self_attn.qkv_proj", + "model.layers.0.self_attn.o_proj", + "model.layers.0.mlp.gate_up_proj", + "model.layers.0.mlp.down_proj", + "lm_head", + ], +) +def test_text_layers_stay_quantized(monkeypatch, prefix): + """The declaration is matched by substring, which is easy to overreach with. + + One declared name that is too short also marks the text tower unquantized, + and vLLM would then build plain float weights for layers the adapter feeds + packed bytes -- losing the whole 30B backbone. + """ + from vllm_gguf_plugin.quantization.utils import is_layer_skipped_gguf + + declared, fused_mapping = _declaration_as_vllm_sees_it(monkeypatch) + + assert not is_layer_skipped_gguf(prefix, declared, fused_mapping) + + +# -------------------------------------------------------------------------- +# Patch embedding +# -------------------------------------------------------------------------- +@pytest.mark.parametrize("patch_temporal", [1, 2, 4]) +def test_patch_embedding_split_adds_back_to_the_stored_sum(patch_temporal): + """Only the sum survives conversion, and still images depend on just the sum. + + The encoder feeds a still image to every time step by expanding the same + patch, so any split that adds back to the stored sum is equivalent for + images -- which is what makes the image path exact rather than approximate. + """ + out_channels = 3 + summed = torch.randn(out_channels, 5) + + blocks = reconstruct_patch_embedding(summed, patch_temporal) + + assert blocks.shape == (out_channels, 5 * patch_temporal) + torch.testing.assert_close( + blocks.reshape(out_channels, patch_temporal, 5).sum(dim=1), summed + ) diff --git a/tests/test_plugin.py b/tests/test_plugin.py index f649df47..5917c535 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -2,7 +2,9 @@ import gc import weakref +from types import SimpleNamespace +import pytest import torch import vllm.engine.arg_utils as arg_utils_module import vllm.model_executor.layers.linear as linear_module @@ -430,3 +432,60 @@ def test_gguf_merged_column_releases_shards_after_concat(monkeypatch): assert layer.qweight is not source_param assert not source_param.data_container assert all(ref() is None for ref in shard_refs) + + +def _supported_mm_limits(quantization: str | None, model_type: str) -> dict: + """Evaluate the patched ``supported_mm_limits`` without starting an engine. + + The wrapper installed by ``register()`` is what is under test, so this calls + the property directly. Reading the effect off a rejected request instead + would mean loading a model, and the assertion would be buried in it. + """ + from vllm.multimodal.processing.context import BaseProcessingInfo + + register() + + hf_config = PretrainedConfig() + hf_config.model_type = model_type + + class Info(BaseProcessingInfo): + def __init__(self): + self.ctx = SimpleNamespace( + model_config=SimpleNamespace( + quantization=quantization, hf_config=hf_config + ) + ) + + def get_supported_mm_limits(self): + return {"image": None, "video": None} + + return dict(Info().supported_mm_limits) + + +def test_gguf_hides_a_modality_the_adapter_cannot_reconstruct(): + """Muse Glimmer's converter keeps only the sum over the patch embedding's + time steps, so video cannot be reconstructed from these weights. + + Dropping the modality here makes vLLM refuse the request with its own + validation error. The alternative is worse than a refusal: video would run + on a reconstruction that is off by roughly 7% in the one channel carrying + frame-to-frame motion, and produce a fluent description of the wrong thing. + """ + assert "video" not in _supported_mm_limits("gguf", "muse_glimmer") + + +def test_the_modality_gate_leaves_images_alone(): + """Positive control: the gate has to be per-modality, not per-model.""" + assert "image" in _supported_mm_limits("gguf", "muse_glimmer") + + +def test_the_modality_gate_is_tied_to_gguf(): + """Negative control: the limitation comes from the GGUF conversion, so + unquantized weights of the same architecture keep video.""" + assert "video" in _supported_mm_limits(None, "muse_glimmer") + + +@pytest.mark.parametrize("model_type", ["gemma3", "qwen2_vl", "llama"]) +def test_the_modality_gate_does_not_reach_other_architectures(model_type): + """Negative control: the patch is global, so it has to stay adapter-scoped.""" + assert "video" in _supported_mm_limits("gguf", model_type) diff --git a/vllm_gguf_plugin/loader.py b/vllm_gguf_plugin/loader.py index ddea77fe..13fdf90e 100644 --- a/vllm_gguf_plugin/loader.py +++ b/vllm_gguf_plugin/loader.py @@ -46,6 +46,7 @@ class GGUFLoadPlan(NamedTuple): files: GGUFModelFiles name_map: dict[str, str] unquantized_modules: tuple[str, ...] + dense_module_suffixes: tuple[str, ...] linear_layouts: dict[str, GGUFLinearLayout] @@ -80,6 +81,25 @@ def _get_unquantized_modules( return tuple(sorted(modules)) +def _publish_declaration_for_a_draft( + model_config: ModelConfig, + plan: GGUFLoadPlan, +) -> None: + """Record the declaration where a draft model will look for it. + + A target model's layers are built against the very ``GGUFConfig`` this + loader extends. A draft's are not: it resolves its own config separately, + rebuilding it from ``hf_config.quantization_config``, so everything recorded + on the shared object reaches the target and never the draft. Writing the + declaration into that dict as well is what lets a draft declare anything. + """ + quant_config = getattr(model_config.hf_config, "quantization_config", None) + if not isinstance(quant_config, dict): + return + quant_config["unquantized_modules"] = list(plan.unquantized_modules) + quant_config["dense_module_suffixes"] = list(plan.dense_module_suffixes) + + class GGUFModelLoader(BaseModelLoader): """ Model loader that can load GGUF files. This is useful for loading models @@ -178,7 +198,11 @@ def _prepare_adapter( ) linear_layouts = adapter.get_linear_layouts(files, model_config, name_map) return adapter, GGUFLoadPlan( - files, name_map, unquantized_modules, linear_layouts + files, + name_map, + unquantized_modules, + tuple(adapter.dense_module_suffixes), + linear_layouts, ) def _iter_weights( @@ -208,6 +232,10 @@ def load_model( logger.debug("GGUF unquantized modules: %s", plan.unquantized_modules) vllm_config.quant_config = cast(GGUFConfig, vllm_config.quant_config) vllm_config.quant_config.unquantized_modules.extend(plan.unquantized_modules) + vllm_config.quant_config.dense_module_suffixes.extend( + plan.dense_module_suffixes + ) + _publish_declaration_for_a_draft(model_config, plan) vllm_config.quant_config.register_linear_layouts( plan.linear_layouts, prefix=prefix, diff --git a/vllm_gguf_plugin/plugin.py b/vllm_gguf_plugin/plugin.py index b76095ad..e8892adc 100644 --- a/vllm_gguf_plugin/plugin.py +++ b/vllm_gguf_plugin/plugin.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -from functools import wraps +from functools import cached_property, wraps from pathlib import Path import vllm.engine.arg_utils as arg_utils_module @@ -48,6 +48,78 @@ def _get_gguf_config_source( return model +def _redirect_draft_to_its_config_source(engine_args) -> str | None: + """Point a separate-file GGUF draft at a directory holding its config. + + The target model gets this for free: ``create_model_config`` rewrites + ``model`` to the config source and keeps the file in ``model_weights``. + A draft has no equivalent, so a ``.gguf`` path reaches ``ModelConfig`` + intact and its config is looked for in the file's own directory -- which, + for a draft shipped next to the target it drafts for, holds the target's + config rather than its own. ``SpeculativeConfig`` then fails validation + before any adapter is consulted. + + Returns the weights path that was redirected away from, so the caller can + put it back on ``model_weights``; ``None`` when there was nothing to do. + """ + speculative_config = engine_args.speculative_config + if not isinstance(speculative_config, dict): + return None + draft_model = speculative_config.get("model") + if not _is_gguf_reference(draft_model): + # The rewrite below edits the caller's dict, so a second pass over the + # same EngineArgs no longer sees a GGUF path. Without the remembered + # value the draft would keep the config directory as its weights source + # and quietly load the unquantized checkpoint sitting there. + return getattr(engine_args, "_gguf_draft_weights", None) + + # Named to match ``EngineArgs.hf_config_path``, which does the same job for + # the target. It has to be removed from the dict either way: the field is + # ours, and ``SpeculativeConfig`` rejects keys it does not declare. + config_path = speculative_config.pop("hf_config_path", None) + source = _get_gguf_config_source(draft_model, None, config_path) + if source == draft_model: + return None + + local = Path(source) + if local.is_dir() and not (local / "config.json").exists(): + raise ValueError( + f"The GGUF speculative draft {draft_model!r} needs a config, and " + f"{source!r} does not contain config.json. GGUF files carry no " + "config.json of their own, and the directory a draft sits in " + "belongs to the model it drafts for. Pass the draft's own config " + 'directory as speculative_config={"hf_config_path": ...}.' + ) + + speculative_config["model"] = source + if speculative_config.get("quantization") is None: + # Without this the draft builds unquantized layers and is then handed + # packed bytes. The target gets the same treatment in + # ``create_model_config``. + speculative_config["quantization"] = "gguf" + engine_args._gguf_draft_weights = draft_model + return draft_model + + +def _mark_draft_config_as_gguf(draft_model_config) -> None: + """Give the draft config the marker ``get_quant_config`` looks for. + + Drafts resolve their quantization config separately from the target, and + that lookup reads ``hf_config.quantization_config`` first and falls back to + ``hf_overrides``. A GGUF file has no ``quantization_config`` to parse, and + the fallback is closed off too: vLLM always hands a draft a *callable* + ``hf_overrides`` so that config transforms applied to the target reach the + draft as well, and the fallback rejects anything that is not a dict. + + The contents do not matter -- ``GGUFConfig.from_config`` ignores them and + the loader fills in the unquantized modules once it has read the file -- + but its presence is what selects that branch. + """ + hf_config = draft_model_config.hf_config + if getattr(hf_config, "quantization_config", None) is None: + hf_config.quantization_config = {"quant_method": "gguf"} + + def _patch_engine_args() -> None: if getattr(EngineArgs, "_gguf_create_model_config_patched", False): return @@ -86,6 +158,8 @@ def create_speculative_config(self, *args, **kwargs): if self.speculative_config is not None: configured_model = configured_model or self.speculative_config.get("model") + draft_weights = _redirect_draft_to_its_config_source(self) + config = original_create_speculative_config(self, *args, **kwargs) gguf_model = self.model_weights if ( @@ -95,6 +169,11 @@ def create_speculative_config(self, *args, **kwargs): and _is_gguf_reference(gguf_model) ): config.draft_model_config.model_weights = gguf_model + if config is not None and draft_weights is not None: + # `model` now names the config directory, so the loader would look + # for weights there and find none. Point it back at the file. + config.draft_model_config.model_weights = draft_weights + _mark_draft_config_as_gguf(config.draft_model_config) return config EngineArgs.create_speculative_config = create_speculative_config @@ -118,6 +197,54 @@ def maybe_override_with_speculators(model, tokenizer, *args, **kwargs): config_module._gguf_speculator_probe_patched = True +def _gguf_unsupported_modalities(model_config) -> tuple[str, ...]: + if getattr(model_config, "quantization", None) != "gguf": + return () + hf_config = getattr(model_config, "hf_config", None) + if hf_config is None: + return () + + from .weights_adapter import get_weights_adapter + + return tuple(get_weights_adapter(hf_config).UNSUPPORTED_MODALITIES) + + +def _patch_mm_limits() -> None: + """Hide modalities an adapter cannot reconstruct from its GGUF weights. + + Wrapping the base ``supported_mm_limits`` rather than each model's + ``get_supported_mm_limits`` covers every subclass at once, and every + consumer -- request validation, the advertised modality list, and profiling + -- reads that one property. Dropping a modality from it makes vLLM reject + those requests with its usual validation error. + """ + from vllm.multimodal.processing.context import BaseProcessingInfo + + if getattr(BaseProcessingInfo, "_gguf_mm_limits_patched", False): + return + + original = BaseProcessingInfo.supported_mm_limits.func + + @wraps(original) + def supported_mm_limits(self): + limits = original(self) + unsupported = _gguf_unsupported_modalities(self.ctx.model_config) + if not unsupported: + return limits + return { + modality: limit + for modality, limit in limits.items() + if modality not in unsupported + } + + patched = cached_property(supported_mm_limits) + BaseProcessingInfo.supported_mm_limits = patched + # Assigning a cached_property after class creation skips __set_name__, which + # is what tells it which attribute to cache under. + patched.__set_name__(BaseProcessingInfo, "supported_mm_limits") + BaseProcessingInfo._gguf_mm_limits_patched = True + + def _register_omni_diffusion_quantization() -> None: try: from vllm_omni.quantization import register_quantization_override @@ -145,4 +272,5 @@ def register() -> None: register_config_parser("gguf")(GGUFConfigParser) _patch_engine_args() _patch_speculator_probe() + _patch_mm_limits() _patch_diffusers_loader() diff --git a/vllm_gguf_plugin/quantization/config.py b/vllm_gguf_plugin/quantization/config.py index 54cef3a8..406e42bf 100644 --- a/vllm_gguf_plugin/quantization/config.py +++ b/vllm_gguf_plugin/quantization/config.py @@ -32,9 +32,20 @@ class GGUFConfig(QuantizationConfig): """Config class for GGUF.""" - def __init__(self, unquantized_modules: list[str] | None = None) -> None: + def __init__( + self, + unquantized_modules: list[str] | None = None, + dense_module_suffixes: list[str] | None = None, + ) -> None: super().__init__() self.unquantized_modules = unquantized_modules or [] + #: Module paths ending in one of these stay unquantized wherever they + #: appear. ``unquantized_modules`` cannot express that: a fused layer + #: is matched by asking whether a declared name *contains* the layer's + #: full runtime path, so every declaration has to spell out a prefix and + #: a layer index. Those are knowable for a target model and awkward for + #: a draft, whose layers vLLM numbers after the target's. + self.dense_module_suffixes = dense_module_suffixes or [] self.linear_layouts: dict[str, GGUFLinearLayout] = {} def __repr__(self) -> str: @@ -56,8 +67,15 @@ def get_config_filenames(cls) -> list[str]: @classmethod def from_config(cls, config: dict[str, Any]) -> "GGUFConfig": - del config - return cls() + # A target model's list is filled in by the loader, which holds the same + # object the layers were built against. A draft's is not: its config is + # rebuilt from scratch here, so anything the loader records reaches the + # target's object and never the draft's. Reading the list back out of + # the config dict is what lets a draft declare one at all. + return cls( + unquantized_modules=config.get("unquantized_modules"), + dense_module_suffixes=config.get("dense_module_suffixes"), + ) @classmethod def override_quantization_method( @@ -76,8 +94,10 @@ def get_quant_method( from .vocal_embeds import GGUFEmbeddingMethod if isinstance(layer, LinearBase): - if is_layer_skipped_gguf( - prefix, self.unquantized_modules, self.packed_modules_mapping + if prefix.endswith(tuple(self.dense_module_suffixes)) or ( + is_layer_skipped_gguf( + prefix, self.unquantized_modules, self.packed_modules_mapping + ) ): return UnquantizedLinearMethod() return GGUFLinearMethod( diff --git a/vllm_gguf_plugin/quantization/params.py b/vllm_gguf_plugin/quantization/params.py index 76c67b24..7cabcaaf 100644 --- a/vllm_gguf_plugin/quantization/params.py +++ b/vllm_gguf_plugin/quantization/params.py @@ -9,6 +9,7 @@ get_tensor_model_parallel_world_size, ) from vllm.model_executor.layers.fused_moe import RoutedExperts +from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.parameter import BasevLLMParameter @@ -19,15 +20,25 @@ def _clone_loaded_weight(loaded_weight: torch.Tensor) -> torch.Tensor: return loaded_weight.detach().clone() +def _gguf_replicated_weight_loader(param, loaded_weight, loaded_shard_id=None): + """Load a whole GGUF weight into an unsharded parameter.""" + param._store(loaded_weight, shard_id=loaded_shard_id) + + def _resolve_gguf_weight_loader( layer: torch.nn.Module, fallback_weight_loader=None, ): - return ( - layer.weight_loader_v2 - if hasattr(layer, "weight_loader_v2") - else fallback_weight_loader - ) + if hasattr(layer, "weight_loader_v2"): + return layer.weight_loader_v2 + if isinstance(layer, ReplicatedLinear): + # ReplicatedLinear is the one linear layer with no v2 loader, and its + # v1 loader asserts that the parameter already has the loaded weight's + # shape. GGUF parameters start out empty and take their shape from the + # packed bytes, so that assertion fires on the first tensor. Nothing is + # sharded here, so storing the tensor whole is the entire job. + return _gguf_replicated_weight_loader + return fallback_weight_loader def _resolve_gguf_weight_type_loader( diff --git a/vllm_gguf_plugin/weights_adapter/__init__.py b/vllm_gguf_plugin/weights_adapter/__init__.py index 715a565c..2f284d16 100644 --- a/vllm_gguf_plugin/weights_adapter/__init__.py +++ b/vllm_gguf_plugin/weights_adapter/__init__.py @@ -11,12 +11,15 @@ get_diffusion_gguf_adapter, ) from .gemma3 import Gemma3GGUFAdapter +from .muse_glimmer import MuseGlimmerDraftGGUFAdapter, MuseGlimmerGGUFAdapter from .olmoe import OLMoEGGUFAdapter from .qwen3_5 import Qwen35GGUFAdapter, Qwen35MtpGGUFAdapter from .transformers import TransformersGGUFWeightsAdapter _ADAPTER_REGISTRY: list[type[BaseGGUFWeightsAdapter]] = [ Gemma3GGUFAdapter, + MuseGlimmerDraftGGUFAdapter, + MuseGlimmerGGUFAdapter, OLMoEGGUFAdapter, Qwen35GGUFAdapter, Qwen35MtpGGUFAdapter, @@ -45,6 +48,8 @@ def get_adapter_architecture(config) -> str | None: "Flux2KleinDiffusionGGUFAdapter", "GGUFModelFiles", "Gemma3GGUFAdapter", + "MuseGlimmerDraftGGUFAdapter", + "MuseGlimmerGGUFAdapter", "OLMoEGGUFAdapter", "QwenImageDiffusionGGUFAdapter", "Qwen35GGUFAdapter", diff --git a/vllm_gguf_plugin/weights_adapter/base.py b/vllm_gguf_plugin/weights_adapter/base.py index b28b1755..ed67d6d5 100644 --- a/vllm_gguf_plugin/weights_adapter/base.py +++ b/vllm_gguf_plugin/weights_adapter/base.py @@ -28,6 +28,21 @@ class BaseGGUFWeightsAdapter(ABC): #: model in speculative decoding) and must stay unquantized. extra_unquantized_modules: tuple[str, ...] = () + #: Suffixes of module paths that must stay unquantized wherever they occur. + #: Use this rather than :attr:`extra_unquantized_modules` when the layers in + #: question are fused and repeated: naming them individually means spelling + #: out a prefix and a layer index, and a draft model's layer indices are + #: assigned relative to the target it drafts for. + dense_module_suffixes: tuple[str, ...] = () + + #: Modalities this adapter cannot reconstruct from GGUF weights. Some + #: converters fold away information that one modality needs while leaving + #: the rest of the model intact; listing it here drops it from the model's + #: supported multimodal limits, so requests carrying it are rejected during + #: input validation instead of running against weights that cannot + #: represent them. + UNSUPPORTED_MODALITIES: tuple[str, ...] = () + @classmethod @abstractmethod def matches(cls, config: PretrainedConfig) -> bool: diff --git a/vllm_gguf_plugin/weights_adapter/muse_glimmer.py b/vllm_gguf_plugin/weights_adapter/muse_glimmer.py new file mode 100644 index 00000000..418ea793 --- /dev/null +++ b/vllm_gguf_plugin/weights_adapter/muse_glimmer.py @@ -0,0 +1,705 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight adapter for Muse Glimmer GGUF checkpoints. + +Two of the conversions this adapter performs are invisible if you get them +wrong: the weights load without complaint and the model generates fluent-looking +but wrong output. Both are called out where they are implemented. + +The first is the Q/K row layout. The conversion script that produces these GGUF +files re-lays out the Q and K projections from the half-split ("NEOX") ordering +that HF checkpoints use into llama.cpp's interleaved ordering. vLLM's Muse +Glimmer implementation hardcodes NEOX rotary embeddings, so the adapter has to +undo that re-layout. + +The re-layout only reorders rows of the output dimension; it never changes a +value. Because GGUF splits quantized super-blocks along the *input* dimension, +every output row is a self-contained run of quantized bytes, so the inverse can +be applied directly to the packed ``qweight`` bytes instead of dequantizing +first. ``test_muse_glimmer_gguf.py`` pins that equivalence bit-for-bit. + +The second is the norm offset, which applies to the per-layer norms but not to +the final one -- see :data:`NORM_OFFSET_SUFFIXES`. +""" + +from __future__ import annotations + +import re +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import torch +from vllm.logger import init_logger +from vllm.model_executor.models.utils import WeightsMapper + +from ..gguf_files import GGUFModelFiles +from ..gguf_utils import maybe_patch_hf_config_from_gguf +from ..weight_utils import get_gguf_tensor_names +from .base import BaseGGUFWeightsAdapter, GGUFWeight + +if TYPE_CHECKING: + from transformers import PretrainedConfig + from vllm.config import ModelConfig + +logger = init_logger(__name__) + +MUSE_GLIMMER_MODEL_TYPES = ("muse_glimmer", "muse_glimmer_text") + +# Neither entry can be looked up in the auto-model mappings: the text-only config +# is in no mapping at all, and while the multimodal one is registered for +# image-text-to-text, the config parser only consults the causal-LM mapping. +# Getting these two crossed does not raise -- pointing ``muse_glimmer`` at the +# causal-LM class quietly builds a text-only model that loads and generates. +MUSE_GLIMMER_ARCHITECTURES = { + "muse_glimmer": "MuseGlimmerForConditionalGeneration", + "muse_glimmer_text": "MuseGlimmerForCausalLM", +} + + +def interleaved_to_neox_row_index( + num_heads: int, + head_dim: int, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Row gather index turning llama.cpp interleaved order back into NEOX. + + The forward direction (what the conversion script did) writes NEOX row ``i`` + to interleaved row ``2i`` for the first half of a head and to ``2i + 1`` for + the second half. Inverting it means gathering the even rows first, then the + odd rows, within each head:: + + [0, 2, 4, ..., head_dim - 2, 1, 3, 5, ..., head_dim - 1] + + This permutation is **not** self-inverse for ``head_dim >= 8``, so applying + the forward helper twice does not undo it. ``head_dim == 4`` happens to be + self-inverse, which makes it a misleading size to test against. + """ + if head_dim % 2: + raise ValueError(f"head_dim must be even, got {head_dim}") + + within_head = torch.cat( + ( + torch.arange(0, head_dim, 2, device=device), + torch.arange(1, head_dim, 2, device=device), + ) + ) + head_offsets = torch.arange(num_heads, device=device) * head_dim + return (head_offsets[:, None] + within_head[None, :]).reshape(-1) + + +def neox_to_interleaved_row_index( + num_heads: int, + head_dim: int, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Row gather index for the forward direction, for tests and debugging.""" + if head_dim % 2: + raise ValueError(f"head_dim must be even, got {head_dim}") + + half = head_dim // 2 + within_head = torch.empty(head_dim, dtype=torch.int64, device=device) + within_head[0::2] = torch.arange(0, half, device=device) + within_head[1::2] = torch.arange(half, head_dim, device=device) + head_offsets = torch.arange(num_heads, device=device) * head_dim + return (head_offsets[:, None] + within_head[None, :]).reshape(-1) + + +def undo_rope_interleave( + tensor: torch.Tensor, + num_heads: int, + head_dim: int, +) -> torch.Tensor: + """Restore NEOX row order on a Q or K tensor. + + Works on packed ``qweight`` bytes (``[num_rows, row_bytes]`` uint8), on + dequantized weights (``[num_rows, input_dim]``) and on 1-D biases alike: + all three only need the leading dimension permuted. + """ + num_rows = tensor.shape[0] + expected = num_heads * head_dim + if num_rows != expected: + raise ValueError( + f"expected {expected} rows for {num_heads} heads of {head_dim}, " + f"got {num_rows}" + ) + + index = interleaved_to_neox_row_index(num_heads, head_dim, tensor.device) + # index_select copies, so the result stays contiguous even though the gather + # is non-monotonic; downstream TP sharding and QKV fusion rely on that. + return tensor.index_select(0, index) + + +TEXT_LAYER_PREFIX = "model.language_model.layers." +VISION_LAYER_PREFIX = "model.vision_tower.layers." +PATCH_EMBEDDING = "model.vision_tower.patch_embedder.patch_embedding.weight" + +# llama.cpp's converter folds the ``1 +`` from this architecture's norm into the +# stored weight, so it has to be taken back out. Only the per-layer norms carry +# the offset; the final norm is stored as-is. Its weights sit close to 1.0, so +# subtracting from it as well would leave values close to 0 and scale the last +# hidden state away entirely. Hence an explicit list of the four per-layer norms +# rather than a test on the ``norm.weight`` suffix, which the final norm matches +# too. +NORM_OFFSET_SUFFIXES = ( + "input_layernorm.weight", + "post_attention_layernorm.weight", + "pre_feedforward_layernorm.weight", + "post_feedforward_layernorm.weight", +) + +# The Q/K norms in these files are not learned parameters: the converter +# synthesizes them as constant tensors from the config's scale factor. vLLM +# applies that factor itself from the config, so loading them would apply it +# twice. Dropping them is what makes the rest of the mapping a bijection onto +# the HF checkpoint. +# +# Skipped before the name map is built rather than mapped to ``None``, so that +# they do not land in the list of tensors reported as unmapped: that list is for +# names the mapping failed to cover, and these are left out on purpose. +SYNTHETIC_QK_NORM_SUBSTRINGS = ("attn_q_norm.", "attn_k_norm.") + +# Modules that have to be handed plain weights rather than packed bytes, for two +# unrelated reasons -- both on the vLLM side, neither fixable from here: +# +# * ``embed_tokens`` is built without forwarding ``quant_config``, so the layer +# only ever owns a plain ``weight`` and rejects packed bytes outright ("no +# module or parameter named model.embed_tokens.qweight_type"). +# * the vision tower's linear layers are either plain ``nn.Linear`` or are also +# built without ``quant_config``, so they have nowhere to put packed bytes. +# +# Unpacking costs roughly 1.9 GB for the embedding and 2.5 GB for the vision +# tower. Both shrink back once vLLM forwards ``quant_config`` in those places. +# +# Every prefix here is also declared unquantized by +# :attr:`MuseGlimmerGGUFAdapter.extra_unquantized_modules`. That is what keeps +# the two decisions from drifting apart: whichever way vLLM builds these layers, +# ``get_quant_method`` sees them as unquantized and expects the plain ``weight`` +# this adapter hands over. Without the declaration, vLLM forwarding +# ``quant_config`` here would make the layer allocate ``qweight`` buffers that no +# incoming tensor matches. +DEQUANTIZED_MODULE_PREFIXES = ( + "model.language_model.embed_tokens", + "model.vision_tower.", + "model.vision_adapter.", + "model.vision_projection", +) + +_VISION_GGUF_PREFIXES = ("v.", "mm.") + + +def has_vision(config: PretrainedConfig) -> bool: + """Whether the vision tower is part of this model. + + Deliberately the same rule vLLM's Muse Glimmer implementation applies, so + that the weights this adapter produces and the modules vLLM builds are + decided by one predicate rather than two. Reading it off the directory + instead -- taking the vision tower to be present whenever a projector file + happens to be resolvable -- decides it a second time, and the two answers + disagree as soon as a config turns vision off. + + The shipped ``config.json`` carries no ``has_vision`` key, so the fallback is + the usual path; the attribute is there for turning vision off explicitly. + """ + configured = getattr(config, "has_vision", None) + if configured is not None: + return bool(configured) + return hasattr(config, "vision_config") + + +def _module_of(param: str) -> str: + return param.rsplit(".", 1)[0] if param.endswith((".weight", ".bias")) else param + + +def dequantized_module_names(mapped_names: Iterable[str]) -> set[str]: + """Full-depth module names for every weight :meth:`transform_weights` unpacks. + + Takes names that have already been mapped into vLLM's namespace, so the rule + can be tested against a handful of synthetic names. + + Declaring only the prefixes above would read as equivalent and is not. + ``is_layer_skipped_gguf`` handles a fused layer by asking whether some + declared name *contains* the full path of each shard, so a prefix -- being + shorter than the path it is a prefix of -- never matches one. The vision + tower's attention is fused into ``qkv_proj``, and these full paths are what + covers it. + """ + return { + _module_of(name) + for name in mapped_names + if name.startswith(DEQUANTIZED_MODULE_PREFIXES) + } + + +def dequantize_packed_rows( + qweight: torch.Tensor, + quant_type: int, + dtype: torch.dtype, + rows_per_chunk: int = 4096, +) -> torch.Tensor: + """Unpack GGUF bytes straight into *dtype*. + + ``gguf.quants.dequantize`` hands back float32, which for a vocabulary this + size is several gigabytes on top of the result. Rows are self-contained -- + GGUF blocks never straddle one -- so converting a block of rows at a time + keeps the transient down to the block rather than the whole tensor. + """ + from gguf import GGML_QUANT_SIZES, GGMLQuantizationType + from gguf.quants import dequantize + + ggml_type = GGMLQuantizationType(quant_type) + block_size, type_size = GGML_QUANT_SIZES[ggml_type] + num_rows, row_bytes = qweight.shape + num_cols = row_bytes // type_size * block_size + + packed = qweight.numpy() + out = torch.empty((num_rows, num_cols), dtype=dtype) + for start in range(0, num_rows, rows_per_chunk): + stop = min(start + rows_per_chunk, num_rows) + values = dequantize(packed[start:stop], ggml_type) + out[start:stop] = torch.from_numpy(values.reshape(stop - start, num_cols)).to( + dtype + ) + return out + + +def reconstruct_patch_embedding( + summed: torch.Tensor, + patch_temporal: int, +) -> torch.Tensor: + """Undo the converter's sum over the patch embedding's time axis. + + The checkpoint keeps one weight block per time step, laid out with time as + the outermost axis of the flattened patch dimension; the converter stores + only their sum, so the individual blocks are gone for good. + + Splitting the sum evenly is exact for still images and is also the closest + available guess at the original blocks. The encoder feeds a still image to + every time step by expanding the same patch, so the output only ever depends + on the sum -- any split reproducing it is equivalent. And the reference + blocks turn out to be near-copies of each other (equal mean magnitude to + within 0.1%, correlation 0.99), so half the sum is close to each of them. + + Video feeds distinct frames per time step, which does depend on the blocks + individually, and is rejected for that reason. + """ + flat = summed.reshape(summed.shape[0], -1) + return flat.div(patch_temporal).repeat(1, patch_temporal) + + +# The vision tower cannot reuse the text substring rules: GGUF calls both +# projections ``attn_q``, while the checkpoint calls the text one +# ``self_attn.q_proj`` and the vision one ``attn.q_proj``. Anchored regexes keep +# the two apart -- they run before the substring pass, so a rewritten vision name +# no longer matches any text rule. +_VISION_BLOCK_LEAVES = { + "attn_q": "attn.q_proj", + "attn_k": "attn.k_proj", + "attn_v": "attn.v_proj", + "attn_out": "attn.proj", + # Straight, not crossed: GGUF ``ffn_up`` is (8960, 1536), which is ``fc1``. + "ffn_up": "mlp.fc1", + "ffn_down": "mlp.fc2", + # These carry no offset, unlike the text norms -- they are stored as-is, so + # they must not appear in NORM_OFFSET_SUFFIXES. + "ln1": "norm1", + "ln2": "norm2", +} + + +def build_muse_glimmer_mapper() -> WeightsMapper: + orig_to_new_regex: dict[re.Pattern, str | None] = { + re.compile(rf"^v\.blk\.(\d+)\.{gguf}\."): (f"{VISION_LAYER_PREFIX}\\1.{hf}.") + for gguf, hf in _VISION_BLOCK_LEAVES.items() + } + orig_to_new_substr: dict[str, str | None] = { + "attn_norm.": "input_layernorm.", + "post_attention_norm.": "post_attention_layernorm.", + "ffn_norm.": "pre_feedforward_layernorm.", + "post_ffw_norm.": "post_feedforward_layernorm.", + "attn_q.": "self_attn.q_proj.", + "attn_k.": "self_attn.k_proj.", + "attn_v.": "self_attn.v_proj.", + "attn_output.": "self_attn.o_proj.", + "attn_gate.": "self_attn.gate_proj.", + "ffn_gate.": "mlp.gate_proj.", + "ffn_up.": "mlp.up_proj.", + "ffn_down.": "mlp.down_proj.", + } + orig_to_new_prefix: dict[str, str | None] = { + "token_embd.": "model.language_model.embed_tokens.", + "blk.": TEXT_LAYER_PREFIX, + "output_norm.": "model.language_model.norm.", + "output.": "lm_head.", + "v.patch_embd.": "model.vision_tower.patch_embedder.patch_embedding.", + "v.position_embd.": ( + "model.vision_tower.patch_embedder.position_embedding_table." + ), + "v.pre_ln.": "model.vision_tower.ln_pre.", + "v.post_ln.": "model.vision_tower.ln_post.", + # Shapes pin these three down: (4096, 6144), (4096, 4096) and + # (6656, 4096) are all distinct and each matches exactly one target. + "mm.0.": "model.vision_adapter.fc1.", + "mm.1.": "model.vision_adapter.fc2.", + "mm.2.": "model.vision_projection.", + } + + return WeightsMapper( + orig_to_new_regex=orig_to_new_regex, + orig_to_new_prefix=orig_to_new_prefix, + orig_to_new_substr=orig_to_new_substr, + ) + + +class MuseGlimmerGGUFAdapter(BaseGGUFWeightsAdapter): + """Adapter for Muse Glimmer GGUF models.""" + + # Only the sum of the patch embedding's per-time-step blocks survives the + # conversion; see reconstruct_patch_embedding. Still images depend on that + # sum alone, so they are exact, while video depends on the blocks + # individually and would run about 7% off in the one channel that carries + # frame-to-frame motion -- wrong in a way that still looks plausible. + UNSUPPORTED_MODALITIES = ("video",) + + def __init__(self) -> None: + self._name_map: dict[str, str] | None = None + + @classmethod + def matches(cls, config) -> bool: + return config.model_type in MUSE_GLIMMER_MODEL_TYPES + + @classmethod + def architecture(cls, config) -> str | None: + return MUSE_GLIMMER_ARCHITECTURES.get(config.model_type) + + def patch_hf_config( + self, + files: GGUFModelFiles, + hf_config: PretrainedConfig, + ) -> PretrainedConfig: + patched = maybe_patch_hf_config_from_gguf( + files.primary_backbone, + hf_config, + mmproj_path=files.mm_proj, + ) + if has_vision(patched) and files.mm_proj is None: + raise RuntimeError( + "The vision tower needs the multimodal projector, and no mmproj " + f"file could be resolved for {files.primary_backbone}. " + "Requesting one quantization does not bring it along, since the " + "projector is quantized separately. Place *mmproj*.gguf beside " + "the backbone, pass model_loader_extra_config={'mm_proj': ...}, " + "or set has_vision=false to run text-only. Continuing without " + "it would leave the vision weights unloaded, which text prompts " + "would not reveal." + ) + return patched + + def build_name_map( + self, + files: GGUFModelFiles, + model_config: ModelConfig, + ) -> dict[str, str]: + mapper = build_muse_glimmer_mapper() + vision = has_vision(model_config.hf_config) + + name_map: dict[str, str] = {} + unmapped: list[str] = [] + for name in sorted(get_gguf_tensor_names(files.all_files)): + if any(substr in name for substr in SYNTHETIC_QK_NORM_SUBSTRINGS): + continue + # The projector is resolved by the loader from the directory, so it + # can turn up even for a config that runs text-only. Leaving its + # tensors out of the map is what keeps them from being loaded. + if not vision and name.startswith(_VISION_GGUF_PREFIXES): + continue + mapped = mapper.apply_list([name]) + if not mapped or mapped[0] == name: + unmapped.append(name) + else: + name_map[name] = mapped[0] + + if unmapped: + logger.warning( + "No HF name for %d Muse Glimmer GGUF tensor(s), skipping: %s", + len(unmapped), + unmapped, + ) + self._name_map = name_map + return name_map + + @property + def extra_unquantized_modules(self) -> tuple[str, ...]: + """Modules this adapter unpacks, which the GGUF stores quantized. + + The loader derives its own list from the tensors a GGUF stores + unquantized, which these are not -- they are quantized in the file and + dequantized here -- so they have to be declared separately. + + Derived from the name map rather than assigned during + :meth:`build_name_map`, so that reading it too early fails loudly. Were + it assigned, a reordering on the loader's side would silently under- + declare, and an under-declared fused layer allocates packed buffers that + no incoming tensor matches. + """ + if self._name_map is None: + raise RuntimeError( + "extra_unquantized_modules was read before build_name_map; the " + "declaration is derived from the name map" + ) + # Both forms are needed. The full paths are what a fused layer matches + # against; the prefixes cover the layers vLLM's own mapper renames on the + # way in (``attn.proj`` becomes ``attn.o_proj``, and a full path recorded + # under the old name stops matching). + return tuple( + sorted( + dequantized_module_names(self._name_map.values()) + | set(DEQUANTIZED_MODULE_PREFIXES) + ) + ) + + def _rope_layout( + self, + name: str, + config: PretrainedConfig, + ) -> tuple[int, int] | None: + """``(num_heads, head_dim)`` to un-interleave *name* with, else ``None``. + + Only Q and K are rotated, so only they were re-laid out. V shares their + shape and quantization, so permuting it as well is not caught by any + shape check -- it just corrupts every value it touches. + + The text and vision towers have to be told apart before the projection + name is inspected, because the text module ``self_attn.q_proj`` also ends + with the vision module's ``attn.q_proj``. + """ + # ``.qweight_type`` is a scalar tag rather than a weight; it shares the + # module prefix, so match on the payload suffixes instead. + if not name.endswith((".qweight", ".weight", ".bias")): + return None + module = name.rsplit(".", 1)[0] + + if name.startswith(TEXT_LAYER_PREFIX): + text_config = config.get_text_config() + if module.endswith("self_attn.q_proj"): + return text_config.num_attention_heads, text_config.head_dim + if module.endswith("self_attn.k_proj"): + return text_config.num_key_value_heads, text_config.head_dim + return None + + if name.startswith(VISION_LAYER_PREFIX): + if not module.endswith(("attn.q_proj", "attn.k_proj")): + return None + # The vision tower is not grouped-query, so Q and K share a count. + vision_config = config.vision_config + num_heads = vision_config.num_attention_heads + return num_heads, vision_config.hidden_size // num_heads + + return None + + def transform_weights( + self, + weights: Iterable[GGUFWeight], + model_config: ModelConfig, + ) -> Iterable[GGUFWeight]: + """Transform mapped GGUF weights into HF-style weights.""" + config = model_config.hf_config + dtype = model_config.dtype + + # The iterator emits a module's ``qweight_type`` immediately before its + # ``qweight``, so a single slot is enough to rejoin the two. + quant_types: dict[str, int] = {} + for name, weight in weights: + module, _, leaf = name.rpartition(".") + if leaf in ("qweight", "qweight_type") and module.startswith( + DEQUANTIZED_MODULE_PREFIXES + ): + if leaf == "qweight_type": + quant_types[module] = int(weight.item()) + continue + name = f"{module}.weight" + weight = dequantize_packed_rows(weight, quant_types.pop(module), dtype) + + if name == PATCH_EMBEDDING: + weight = reconstruct_patch_embedding( + weight, config.vision_config.patch_temporal + ) + elif name.endswith(NORM_OFFSET_SUFFIXES): + weight = weight - 1 + elif (layout := self._rope_layout(name, config)) is not None: + weight = undo_rope_interleave(weight, *layout) + yield name, weight + + +MUSE_GLIMMER_DRAFT_MODEL_TYPES = ("muse_glimmer_assistant",) + +# vLLM serves this draft with the generic DFlash head, and EAGLEConfig rewrites +# the architecture to DFlash{arch} on the way in, so the registry is asked for +# DFlashMuseGlimmerAssistantModel. The bare name is what belongs here: the +# rewrite happens after the config parser runs. +MUSE_GLIMMER_DRAFT_ARCHITECTURE = "MuseGlimmerAssistantModel" + +# Layer-local renames. The draft is a plain requantization -- see +# MuseGlimmerDraftGGUFAdapter -- so this table is the whole adapter. +_DRAFT_WITHIN_LAYER = { + "attn_norm": "input_layernorm", + "ffn_norm": "post_attention_layernorm", + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", + "attn_q_norm": "self_attn.q_norm", + "attn_k_norm": "self_attn.k_norm", + "ffn_gate": "mlp.gate_proj", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", +} + +# The three tensors outside the blocks. ``fc`` consumes the concatenated +# hidden states of the target layers named by ``target_layer_ids``, which is why +# its input dimension is a multiple of the hidden size. +_DRAFT_TOP_LEVEL = { + "fc.weight": "encoder.fc.weight", + "enc.output_norm.weight": "encoder.output_norm_enc.weight", + "output_norm.weight": "norm.weight", +} + +_DRAFT_BLOCK_RE = re.compile(r"^blk\.(\d+)\.(.+)\.weight$") + +# The projections that make up the fused ``qkv_proj``. See +# MuseGlimmerDraftGGUFAdapter.dense_module_suffixes for why these three, and +# only these three, cannot stay packed. +_DRAFT_DENSE_SUFFIXES = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", +) + + +def _draft_model_type(config: PretrainedConfig) -> str | None: + """Read the draft's own model type, seeing through EAGLEConfig. + + A dflash draft is wrapped in ``EAGLEConfig`` on its way into the engine, + which reports ``model_type == "eagle"`` and keeps the real config on + ``.model``. The wrapper is applied after the config parser has run, so + ``architecture()`` sees the bare type while the loader sees the wrapped + one, and an adapter that only checks the bare type stops matching exactly + when the weights are about to be mapped -- the fallback adapter then takes + over and fails on an architecture it has never heard of. + """ + model_type = getattr(config, "model_type", None) + if model_type != "eagle": + return model_type + return getattr(getattr(config, "model", None), "model_type", None) + + +def build_muse_glimmer_draft_name_map(tensor_names: Iterable[str]) -> dict[str, str]: + """Map the draft's GGUF tensor names to assistant-checkpoint names.""" + name_map: dict[str, str] = {} + unmapped: list[str] = [] + for name in sorted(tensor_names): + if (mapped := _DRAFT_TOP_LEVEL.get(name)) is not None: + name_map[name] = mapped + continue + match = _DRAFT_BLOCK_RE.match(name) + within = _DRAFT_WITHIN_LAYER.get(match.group(2)) if match else None + if within is None: + unmapped.append(name) + continue + name_map[name] = f"layers.{match.group(1)}.{within}.weight" + + if unmapped: + logger.warning( + "No HF name for %d Muse Glimmer draft tensor(s), skipping: %s", + len(unmapped), + unmapped, + ) + return name_map + + +class MuseGlimmerDraftGGUFAdapter(BaseGGUFWeightsAdapter): + """Adapter for the Muse Glimmer DFlash draft, shipped as its own GGUF. + + None of the conversions the backbone needs apply here, and that is the + point worth stating: the draft is a plain requantization. Its Q/K rows are + already in NEOX order, its norms are stored as-is, and its Q/K norms are + learned rather than synthesized. Reusing the backbone's rules would rewrite + correct weights -- the draft would still load and still produce fluent text, + because the target verifies every token, but its proposals would stop being + accepted. ``test_muse_glimmer_dflash_gguf.py`` pins each of those three + facts. + + So the name map is the whole conversion. The one thing + ``transform_weights`` does is unpack Q/K/V, which the fused KV path forces + and which changes no values. + """ + + #: Never present in the draft's own checkpoint -- both are shared from the + #: target after loading, so they must stay ordinary vocab modules rather + #: than ones expecting packed GGUF bytes. + extra_unquantized_modules = ("embed_tokens", "lm_head") + + #: The head fuses every layer's KV projection into one buffer at the end of + #: loading and reads ``qkv_proj.weight`` to build it, which a quantized + #: layer does not have. Leaving these packed is therefore not an option, + #: and unpacking them is cheap: Q/K/V are about a sixth of the draft, so it + #: still loads at roughly a third of its unquantized size. Everything else + #: stays quantized. + dense_module_suffixes = ("self_attn.qkv_proj",) + + def __init__(self) -> None: + self._dense_modules: tuple[str, ...] = () + + @classmethod + def matches(cls, config) -> bool: + return _draft_model_type(config) in MUSE_GLIMMER_DRAFT_MODEL_TYPES + + @classmethod + def architecture(cls, config) -> str | None: + if not cls.matches(config): + return None + return MUSE_GLIMMER_DRAFT_ARCHITECTURE + + def build_name_map( + self, + files: GGUFModelFiles, + model_config: ModelConfig, + ) -> dict[str, str]: + del model_config + name_map = build_muse_glimmer_draft_name_map( + get_gguf_tensor_names(files.all_files) + ) + self._dense_modules = tuple( + sorted( + name.removesuffix(".weight") + for name in name_map.values() + if name.endswith(_DRAFT_DENSE_SUFFIXES) + ) + ) + return name_map + + def transform_weights( + self, + weights: Iterable[GGUFWeight], + model_config: ModelConfig, + ) -> Iterable[GGUFWeight]: + """Unpack the Q/K/V projections; pass everything else through. + + No values change here. Dequantizing is a representation change that the + fused KV path forces, not one of the conversions the backbone needs -- + see the class docstring for why none of those apply to the draft. + """ + dtype = model_config.dtype + dense = set(self._dense_modules) + + # The iterator emits a module's ``qweight_type`` immediately before its + # ``qweight``, so a single slot is enough to rejoin the two. + quant_types: dict[str, int] = {} + for name, weight in weights: + module, _, leaf = name.rpartition(".") + if leaf in ("qweight", "qweight_type") and module in dense: + if leaf == "qweight_type": + quant_types[module] = int(weight.item()) + continue + name = f"{module}.weight" + weight = dequantize_packed_rows(weight, quant_types.pop(module), dtype) + yield name, weight