[WIP] Shared-input Hessian dedup: :in=<tag> plan metadata + CPU probe + looper leader/follower capture

[Qubitium]

Summary

Model-definition metadata that identifies decoder-layer modules consuming the same input activation, and looper wiring so the GPTQ Hessian (H = XᵀX) is collected once per group and copied to the other members (e.g. q/k/v_proj, gate/up_proj).

Phase 1 (8bf42b7, e7cfd32): metadata + CPU probe. Phase 2 (e7ec666): looper dedup.

Contract (after review): dedup is explicit opt-in only. Every quantizable module is a singleton group unless sibling leaves under the same parent carry the same :in=<tag>. Subset digits are never used to infer sharing.

What Changed

• New gptqmodel/models/shared_input.py:
  • :in=<tag> module_tree leaf flag (ordinary flag to existing parsers; does not alter emitted blocks).
  • build_shared_input_plan(module_tree, layer_modules) -> SharedInputPlan with frozen SharedInputGroup(key, parent, modules, subset_indices, explicit); helpers leader_for, followers_of, shares_input, is_explicit, for_subset, filter_modules, with_prefix, dedup_count.
  • Grouping: untagged -> singleton (key == module path); same parent + same :in=<tag> -> shared (key == "<parent>:in=<tag>", may span subsets). :!/:? leaves excluded; experts.{i} expand per expert (routed experts never share with each other).
  • collect_leaf_specs() raises ValueError when a leaf is repeated across module_tree variants with different subset_tag / input_tag / quantize; identical duplicates are fine.
  • probe_shared_inputs(layer, plan, forward) -> SharedInputProbeReport: pre-hooks every planned module, runs a real forward, compares captured inputs (shape/dtype/values/call-count). Status: has_errors (mismatches or undeclared identical inputs), fully_verified (no errors, no missing modules, no uncalled groups), ok == fully_verified (strict).
• BaseQModel.shared_input_plan(model_config, quantize_config, is_awq_quantize=False) classmethod.
• :in=x tags added only to definitions covered by the tiny-CPU real-forward suite: llama (+ qwen2/mistral/gemma2/gemma3 via inheritance), qwen3, phi3, qwen3_moe, qwen2_moe, mixtral, deepseek_v3 (q_a_proj/kv_a_proj_with_mqa -> in=h), glm4_moe (incl. nested shared_experts spanning blocks), qwen3_5 / qwen3_5_text / qwen3_5_moe / qwen3_5_moe_text (incl. linear_attn.in_proj_qkv:0 + in_proj_z:1), qwen3_next, gpt_oss, llama4.
• MLA definitions (deepseek_v2/v3/v32/v4/vl_v2, glm4_moe_lite, glm5_next, glm_moe_dsa, kimi_k25, longcat_flash, minicpm3, axk2) keep q_b_proj:…:in=q_a, kv_b_proj:…:in=kv_a as documented singletons — the probe showed they receive different tensors ((2,6,16) vs (2,1,6,16)).
• README section under "How to Add Support for a New Model".
• full_layer_modules()/simple_layer_modules() output unchanged for all definitions (:in= is transparent to block emission).

Phase 2 — looper Hessian dedup

• HessianConfig.dedup_shared_inputs: bool = True (serialized to quant meta; dynamic-overridable per module, e.g. {".*\\.v_proj$": {"hessian": {"dedup_shared_inputs": False}}}).
• GPTQ.adopt_hessian_from(leader): leader.materialize_global_hessian() then a private fp32 copy of H (copy=True, follower's target device) plus nsamples/fwd_counter; clears follower partials. Raises on column mismatch. Independence matters because quantize() mutates H in place and free() drops it.
• LoopProcessor.begin_shared_input_capture(model, subset_names, is_lm_head_module) -> {follower: leader} / end_shared_input_capture(subset_names): no-op defaults; GPTQProcessor overrides them:
  • plan derived once per model class via model.shared_input_plan(model_config, quantize_config);
  • per subset pass, for each explicit group: candidates = members present in this subset, with a task of type exactly GPTQ (GPTAQ/FOEM excluded), per-module hessian.dedup_shared_inputs on, same columns; first in subset order leads, others follow; <2 candidates -> no dedup;
  • pre_process_fwd_hook(name) returns a no-op hook for followers (leader/singleton hooks unchanged, incl. keep-mask splitting);
  • end_shared_input_capture -> follower.adopt_hessian_from(leader), marks follower as having captured input ids, clears election state; shared_input_dedup_count tracks adoptions.
• stage_subset._run_single_subset_pass: begin_… before hooks are installed (only when execute_forward), end_… right after hooks are removed and before coverage validation / worker quantization.
• Not touched: AWQ/QQQ/ParoQuant processors, MoE routing/early-stop, lm_head (never dedups), LoopProcessor lifecycle.
• Follow-up: opting more definitions in once they have real-forward coverage.

Tests

• I added a new simple/fast unit test for this change, or documented why that is not applicable.
• I ran the new targeted test locally before opening this PR.
• I ran any other directly relevant local tests.

tests/module_tree/test_shared_input.py: flag parsing, leaf-spec extraction incl. repeated nested keys, expert placeholder resolution, singleton default, explicit same-subset / cross-subset grouping, conflicting-variant metadata raises (tag, :!, :?, subset) while identical duplicates pass, :!/:? exclusion, block output unchanged with :in=, prefix/filter, concurrent determinism, MoE, every registered definition covers its quantizable paths and has no non-explicit shared group, synthetic probe pass/fail modes (shape, dtype, value, missing, uncalled, partial, multi-call, kwarg tensors, hook cleanup on success/exception, undeclared pairs, has_errors/fully_verified/strict ok).

tests/module_tree/test_shared_input_cpu_forward.py: tiny HF configs on CPU (hidden 32–64, 2 layers) through convert_model + real model(input_ids) probe for Llama, Qwen2/3, Mistral, Gemma2/3, Phi3, Qwen2/3-MoE, Mixtral, DeepSeek-V3, GLM4-MoE, Qwen3.5 dense + MoE, Qwen3-Next, GPT-OSS, Llama4. Every layer must be fully_verified with undeclared == () and every shared group explicit; negative tests show a wrong MLA tag -> mismatches, an untagged q/k/v triple -> undeclared, a wrongly split q/k -> undeclared, and un-routed experts -> unverified (strict ok False, has_errors False).

tests/test_shared_input_hessian_dedup.py (26, CPU): adopt_hessian_from equal-but-independent H (mutating follower leaves leader intact), nsamples/fwd_counter carry-over, idempotent multi-follower, uncalled leader -> zero H, self no-op, column mismatch raises, quantize() after adopt == independent collection; election: first-in-subset leads, singletons/untagged never, subset-restricted members, missing tasks, dedup_shared_inputs=False (global + dynamic), lm_head, model without plan API, non-plain-GPTQ excluded, column mismatch skipped, state reset on next begin; capture: followers get 0 add_batch then adopt, hooks resume normal after end, dedup vs independent H allclose, keep-mask path, end without begin / with pruned tasks.

tests/test_shared_input_hessian_dedup_e2e.py: real GPTQModel.load + quantize() of a 2-layer tiny Llama (GQA, hidden 64) on CPU with dedup on vs off: elections are {k,v -> q} / {up -> gate} per layer, followers' fwd_counter == 0, and all 14 quantized weight tensors are bit-identical between the two runs.

cd tests && python -m pytest test_shared_input_hessian_dedup.py test_shared_input_hessian_dedup_e2e.py -q
# 27 passed in 3.3s
cd tests && python -m pytest module_tree/test_shared_input.py module_tree/test_shared_input_cpu_forward.py module_tree/test_module_tree_variants.py -q
# 248 passed in 4.17s
cd tests && python -m pytest module_tree -q --deselect module_tree/test_moe_flag_parsing.py::TestMoEModuleName::test_get_moe_module_name_none_tree --deselect module_tree/test_subset.py::test_qwen3_5_moe_subset_early_stop_follows_module_tree_execution_order
# 333 passed, 14 skipped
cd format && ruff check --config ruff.toml ../gptqmodel/models/shared_input.py ../gptqmodel/models/base.py ../gptqmodel/models/definitions/*.py ../tests/module_tree/test_shared_input*.py
# All checks passed!

Pre-existing on main (unrelated, also fail without this branch): test_moe_flag_parsing.py::TestMoEModuleName::test_get_moe_module_name_none_tree, test_subset.py::test_qwen3_5_moe_subset_early_stop… (needs CUDA).

Review Requirements

• I personally reviewed every file in this diff.
• I checked that the code matches existing project structure, APIs, and conventions.
• I avoided unnecessary monkeypatching and used the project's normal extension points where possible.

Notes

WIP/draft. Review round 1 (comments 5536991865 / 5536994939 / 5536996271) addressed in e7cfd32: singleton default + explicit opt-in, conflict detection across variants, strict ok. Qwen3_5_MoeQModel / Llama4QModel (vision-language wrappers) are tagged alongside their text-only classes because they share the same decoder layer module; only the text variants are exercised by the CPU suite.

Link to Devin session: https://app.devin.ai/sessions/62474a853aa048bf95738e1e03f06e20
Open in Devin Desktop: https://app.devin.ai/desktop/session/62474a853aa048bf95738e1e03f06e20?variant=devin
Requested by: @Qubitium

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[Qubitium]

Review: P1 — do not infer shared-Hessian eligibility from subset grouping by default

The current default rule — same immediate parent + same expanded quantization subset/block => same input tensor — is not a safe semantic contract.

The numeric subset/block historically describes quantization/execution grouping, not tensor identity. This PR itself demonstrates the problem: MLA definitions had modules in the same subset that receive different activations, so they needed explicit :in= splits. That is evidence that subset membership is not sufficient proof of shared input.

The risk becomes correctness-critical in the follow-up that skips follower add_batch: a false-positive group will silently reuse the wrong H = XᵀX for another module and corrupt the quantization result without necessarily crashing.

The real-forward suite is good, but CASES currently validates only a relatively small subset of the registered model definitions. The all-definition test checks structural coverage/consistency; it does not prove that the inferred tensors are identical at runtime.

I recommend one of these safer contracts before wiring this into the looper:

1. Safest: default every module to a singleton and require explicit :in=<tag> opt-in for every dedup group.
2. Or add an explicit per-definition opt-in such as infer_shared_inputs_from_subsets = True, enabled only for definitions with real-forward verification.
3. At minimum, do not consume inferred/default groups for Hessian dedup until every dedup-enabled definition has runtime tensor-identity coverage.

The optimization is valuable, but the failure mode here is silent quantization error, so the default should be conservative rather than inferred from an unrelated grouping primitive.

[Qubitium]

Review: P1 — collect_leaf_specs() silently discards conflicting metadata across module-tree variants

collect_leaf_specs() currently does:

if template in specs:
    return

so the first definition wins for a repeated relative leaf path across module_tree variants.

That is unsafe for :in= because BaseQModel.build_layer_modules() intentionally supports/merges multiple complete tree variants. If two variants contain the same relative module path but assign different :in= tags (or different quantize/subset metadata), the planner silently uses whichever variant was visited first.

The docstring even says the first definition wins "so variant trees cannot silently re-tag a shared leaf," but this implementation does the opposite: it silently hides the conflict.

Before this metadata can drive Hessian reuse, repeated leaf specs should be validated. If an existing spec differs in any semantic field (input_tag, quantize, and probably the relevant grouping metadata), raise a clear error rather than accepting the first value. Identical duplicate specs are fine.

Suggested shape:

new = LeafSpec(...)
old = specs.get(template)
if old is not None and old != new:
    raise ValueError(f"conflicting module_tree metadata for {template}: {old} vs {new}")
specs.setdefault(template, new)

A unit test with two module-tree variants that repeat the same leaf with different :in= tags would catch this.

[Qubitium]

Review: P2 — SharedInputProbeReport.ok can report success when the plan was not actually verified

ok is currently:

return not self.mismatches and not self.undeclared

so a probe with missing planned modules and/or entirely uncalled groups can still return True. The tests explicitly lock this behavior in (test_missing_and_uncalled_modules, test_moe_unrouted_expert_is_reported_unverified).

That makes report.ok easy to misuse as a validation gate. For a feature intended to prove Hessian-sharing safety, "no contradiction observed" is different from "verified." A missing module or uncalled group provides no evidence that its declared/inferred sharing is correct.

I would separate these states explicitly, for example:

• has_errors: mismatches / undeclared conflicts
• fully_verified: no errors and no missing modules and no unverified groups
• keep ok only if its meaning is documented as "no observed contradiction"

Or make ok strict and add a separate permissive property for MoE probes where unrouted experts are expected.

This matters because future CI or tooling is very likely to write only assert report.ok; with the current API that can silently bless an incompletely exercised plan.

[Qubitium]

Review summary

I reviewed the planner/probe implementation, the model-definition annotations, the targeted synthetic tests, the tiny real-forward coverage, existing PR discussion, and current CI state.

Verdict: changes requested before this becomes a Hessian-dedup contract. The implementation is well-tested structurally, but there are two correctness hazards that can turn into silent quantization error once follower Hessian collection is skipped:

1. Defaulting to same parent + same subset => same input relies on subset metadata that was not originally an input-identity guarantee, while only a subset of model definitions receive real-forward validation.
2. Conflicting :in= metadata across module-tree variants is silently resolved by first-definition-wins.

I also flagged the probe API semantics because report.ok can be true for missing/unverified groups, which is dangerous for future CI validation.

Current Ruff workflow is green; I did not find a CI failure driving these findings. The concerns are semantic/correctness issues in the new contract rather than formatting or test hygiene.

[devin-ai-integration]

Re "P1 — do not infer shared-Hessian eligibility from subset grouping by default" (comment 5536991865): agreed, fixed in e7cfd32. Untagged modules are now singleton groups (key = module path); dedup only happens for same-parent leaves carrying the same :in=<tag>, and subset digits are never consulted for grouping. :in= tags were added only to definitions exercised by the tiny-CPU forward suite (llama family, qwen3, phi3, qwen2/qwen3 MoE, mixtral, deepseek v3, glm4 MoE, qwen3.5 dense/MoE, qwen3-next, gpt-oss, llama4). That suite now asserts every multi-module group is explicit, report.fully_verified, and undeclared == (), so both a wrong tag and a missing opt-in fail the test instead of silently merging. Definitions without runtime coverage (incl. the MLA q_b/kv_b singletons) get no dedup.

[devin-ai-integration]

Re "P1 — collect_leaf_specs() silently discards conflicting metadata" (comment 5536994939): fixed in e7cfd32. record() now builds the new LeafSpec and raises ValueError("conflicting module_tree metadata for across variants: <old> vs <new>") when an existing spec differs in any field (subset_tag, input_tag, quantize); identical duplicates are still accepted. Tests: test_conflicting_metadata_across_variants_raises, test_conflicting_quantize_or_subset_flags_raise (parametrized over :!, :?, subset, tag drop) and test_identical_metadata_across_variants_is_allowed; the all-definitions test confirms no registered variant tree currently conflicts.

[devin-ai-integration]

Re "SharedInputProbeReport.ok is permissive" (comment 5536996271): fixed in e7cfd32. The report now exposes has_errors (mismatches or undeclared), fully_verified (no errors, no missing_modules, no unverified), and ok is a strict alias of fully_verified. The real-forward suite gates on ok; the un-routed-expert test now asserts not has_errors and not ok instead of ok, and new synthetic tests cover missing-module-only and uncalled-group-only cases both failing strict ok.

[Qubitium]

@ZX-ModelCloud Validate

[ZX-ModelCloud]

Follow-up validation review

I reviewed the current head (e7ec6669730e429401cf0a2c4bcc9623dd82ba3d) and recommend changes requested before merge.

P1 — LlamaQModel tags implicitly enable dedup for unverified subclasses

The conservative contract introduced in this PR says that only model definitions verified through a real forward should carry :in= tags. However, the tags were added directly to LlamaQModel.module_tree, while HessianConfig.dedup_shared_inputs defaults to True.

Those tags are inherited automatically by definitions such as DeciLM, Dream, Ernie 4.5, InternLM, Instella, MobileLLM, Xverse, and other Llama-derived classes that do not override module_tree. Enumerating MODEL_MAP at this head shows 41 model-type entries with shared groups, while the real-forward CASES suite contains only 17 cases. Therefore, the statement that tags were added only to definitions exercised by the CPU suite is not currently true.

I additionally ran real-forward probes for Cohere, Cohere2, Gemma, Granite, OLMo2, and StableLM; all passed. That still leaves the trust-remote-code descendants and other inherited definitions unverified. Because an incorrect inherited tag silently skips follower Hessian capture and reuses the wrong H, I recommend either:

1. moving the opt-in metadata to each verified concrete definition, leaving unverified descendants disabled; or
2. adding real-forward coverage for every inherited model type that receives shared groups.

Relevant locations:

• gptqmodel/models/definitions/llama.py:17-26
• gptqmodel/quantization/config.py:1214-1220
• tests/module_tree/test_shared_input_cpu_forward.py:64-186

P2 — cross-subset groups are not deduplicated, and dedup_count overstates the result

The planner and README allow :in= groups to span subsets, including Qwen3.5's linear_attn.in_proj_qkv:0 and linear_attn.in_proj_z:1. However, GPTQProcessor.begin_shared_input_capture() elects members only from the current subset_names and skips groups with fewer than two members in that pass. Consequently, a group with one member in each subset performs no Hessian dedup at runtime.

SharedInputPlan.dedup_count still counts every follower across the full group, so it reports a collection as skippable even when the looper will skip none. Either retain/reuse the leader Hessian across subset passes, or scope the API, count, metadata, and documentation to same-subset dedup.

Relevant locations:

• gptqmodel/looper/gptq_processor.py:245-268
• gptqmodel/models/shared_input.py:227-229
• gptqmodel/models/definitions/qwen3_5.py:41-45

Validation performed

• Exact PR head: e7ec6669730e429401cf0a2c4bcc9623dd82ba3d
• Targeted suite: 275 passed in 12.42s
• git diff --check: clean
• All currently visible GitHub checks are green

The three earlier review issues are resolved correctly: untagged leaves now remain singletons, conflicting variant metadata raises, and report.ok is a strict alias of fully_verified. The same-subset Hessian adoption path is otherwise well structured: followers receive independent fp32 copies, so later in-place quantization and free() do not alias the leader's Hessian.