GLM-4 GPTQ import: encoder remap + fused-QKV split - #424
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
This PR improves GLM-4 / ChatGLM compatibility in Mobius by normalizing legacy HuggingFace config fields and adding GPTQ checkpoint import handling for the original transformer.encoder.* weight layout, including splitting fused QKV projections after GPTQ preprocessing while keeping fused gate/up projections intact.
Changes:
- Update
ChatGLMCausalLMModelpreprocessing to remap legacy checkpoint key prefixes and split fused QKV tensors intoq_proj/k_proj/v_projweights after GPTQ conversion. - Extend
ArchitectureConfig.from_transformers()to recognize additional ChatGLM/GLM-4 legacy config fields (e.g.,num_layers,kv_channels,seq_length, multi-query group fields, bias flags) and set ChatGLM-specific RoPE defaults. - Add focused unit tests for config extraction and GPTQ fused-QKV import/splitting.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/mobius/models/chatglm.py |
Remaps legacy transformer.encoder.* keys to Mobius canonical names, runs generic GPTQ preprocessing, then splits fused QKV into separate projections while using fused gate/up import. |
src/mobius/models/chatglm_test.py |
Adds a targeted test covering legacy-name remap + GPTQ preprocessing + fused QKV split behavior. |
src/mobius/_configs/_base.py |
Enhances config normalization for ChatGLM/GLM-4 legacy fields (layers/head_dim/KV heads/bias flags/seq_length) and sets partial_rotary_factor=0.5 for ChatGLM implicit RoPE defaults. |
src/mobius/_configs_test.py |
Adds a regression test for ChatGLM legacy config-field extraction. |
| # LLaDA/OLMo expose the activation as ``activation_type`` (e.g. "silu"). | ||
| or getattr(config, "activation_type", None) | ||
| or ("silu" if model_type in ("qwen",) else None) | ||
| or ("silu" if model_type in ("qwen", "chatglm") else None) |
There was a problem hiding this comment.
Updated the _resolve_hidden_act() fallback documentation to cover both Qwen and ChatGLM SiLU fallbacks.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
src/mobius/_configs/_base.py:69
- The
_resolve_hidden_act()docstring still documents the hardcoded SiLU fallback as Qwen-only, but the implementation now also falls back to SiLU formodel_type == "chatglm". Update the docstring so it stays accurate for future readers.
# LLaDA/OLMo expose the activation as ``activation_type`` (e.g. "silu").
or getattr(config, "activation_type", None)
or ("silu" if model_type in ("qwen", "chatglm") else None)
# gelu_activation is a boolean (XLM) — must be after all string
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Added P0 MLA attention-lowering guard in commit
|
| @pytest.mark.arch_validation | ||
| def test_deepseek_v2_lite_unequal_kv_heads_retain_attention(self): | ||
| """CUDA export must not rewrite unequal-dimension MLA K/V to GQA.""" | ||
| with pytest.warns(UserWarning, match="GQA fusion expected"): | ||
| pkg = build( |
There was a problem hiding this comment.
Replaced the CUDA/HF arch-validation case with the self-contained CPU test in src/mobius/rewrite_rules/_group_query_attention_test.py. Its unequal and equal MLA K/V cases are not marked arch_validation; both are collected by the main CI selection and were run locally.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
P0.2 expert-major int4 QMoE exporter landed in commit
|
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Copilot review follow-ups are resolved in
Validation: CI-selected GQA suite 24 passed / 1 deselected; ChatGLM tests 2 passed; Ruff and |
| def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): | ||
| """Return selection and aggregation tensors for QMoE.""" | ||
| weight_t = op.Transpose(self.weight, perm=[1, 0]) | ||
| router_logits = op.MatMul(hidden_states, weight_t) | ||
| router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) | ||
| return router_logits, None, True, 1.0 |
There was a problem hiding this comment.
Verified against the ORT QMoE kernels — passing raw logits here is intentional and correct, so no routing-math change was needed (fixed by adding a clarifying docstring in 5b0903c).
TopKGate.qmoe_routing returns (router_logits_f32, router_weights=None, normalize=True, 1.0). With router_weights omitted and normalize_routing_weights=1, com.microsoft::QMoE reproduces TopKGate.forward exactly:
- Selection: top-k is taken over
router_probsby value. Softmax is monotonic, so argmax over logits == argmax over softmax — selection is identical. - Aggregation (router_weights=None path): the kernel computes softmax over the k selected logits. See
contrib_ops/cpu/moe/moe_quantization_cpu.cc(default branch):route_scale = exp(logit_j - max) / Σ_{top-k} exp(logit). The CUDA path (ft_moe/moe_kernel.cu, fusedtopk_gating_softmax/moe_softmax+moe_top_k) computes full softmax then, fornormalize=1, renormalizes over the selected k — the full-softmax denominator Z cancels, giving the sameexp(logit)/Σ_{top-k} exp(logit).
That is precisely Softmax(TopK(router_logits)), i.e. TopKGate.forward.
Re: why router_weights differs across gates — the ORT schema (contrib_defs.cc, QMoE Input 14) says when router_weights is omitted, router_probs is used for both selection and the internal softmax-over-selected; when provided, its values are gathered at the selected experts (no internal softmax) and optionally renormalized. So TopKGate correctly delegates the softmax-over-selected to the op via None, whereas SoftmaxTopKGate/SigmoidTopKGate must pass their already-activated probs as router_weights to avoid a double softmax (or softmax-of-sigmoid). All three are mathematically consistent with their forward references.
There was a problem hiding this comment.
Correction to the prior reply: ORT CPU QMoE honors Input 14 (router_weights) and gathers it at selected experts, but CUDA QMoE ignores Input 14 and always runs softmax-top-k on router_probs (contrib_ops/cpu/moe/moe_quantization_cpu.cc; contrib_ops/cuda/moe/moe_quantization.cc). Therefore raw logits with router_weights=None in TopKGate are correct on both EPs: selected experts match by monotonic ordering and their weights equal Softmax(TopK(logits)). Activated router_weights is CPU-correct but double-activated on CUDA.
| from transformers import AutoConfig | ||
|
|
||
| hf_config = AutoConfig.from_pretrained( | ||
| "deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True | ||
| ) |
There was a problem hiding this comment.
Fixed in 5b0903c: the test is now self-contained and the arch_validation marker is removed, so it runs in the offline main CI (pytest -m 'not integration and not arch_validation').
Instead of downloading deepseek-ai/DeepSeek-V2-Lite from HuggingFace, it builds a tiny DeepSeek-V2-Lite-shaped model from an inline ArchitectureConfig (MLA attention, first_k_dense_replace=1, softmax routing, shared experts). hidden_size/moe_intermediate_size are multiples of the quant group_size (128) as QMoE requires; the small MLA/dense dims are fine for MatMulNBits (ceil-padded K). No network, no CUDA required.
The invariant is unchanged and still asserted: one com.microsoft::QMoE per routed MoE layer (2 QMoE / 3 Attention here), shared experts stay dense MatMulNBits, QMoE float routing inputs (1/3/6/14), and no .moe.experts. initializers leak into the graph. Verified green offline: pytest src/mobius/rewrite_rules/_group_query_attention_test.py -k deepseek_v2_lite -v passes with HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1.
Address Copilot review feedback questioning whether TopKGate.qmoe_routing should pass raw router logits (instead of probabilities) as the QMoE `router_probs` input. This is intentional and correct: with `router_weights=None` and `normalize_routing_weights=1`, com.microsoft::QMoE reproduces `TopKGate.forward` exactly — it selects the top-k experts by logit value (softmax is monotonic) and, because `router_weights` is omitted, computes aggregation weights as softmax over the k selected logits, i.e. `Softmax(TopK(router_logits))`. Verified against the ORT QMoE kernels: the CPU default path in contrib_ops/cpu/moe/moe_quantization_cpu.cc computes `exp(logit - max) / sum` over the selected experts, and the CUDA fused/unfused paths in contrib_ops/cuda/moe/ft_moe/moe_kernel.cu apply full softmax then renormalize over the selected k when normalize=1 (the softmax denominator cancels). The Softmax/Sigmoid gates instead pass their activated probabilities via `router_weights` to bypass the op's internal softmax-over-selected. No routing math changes; adds a clarifying docstring only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
`test_deepseek_v2_lite_int4_uses_one_qmoe_per_moe_layer` was marked
`@pytest.mark.arch_validation` but lives under `src/`, so it ran in no CI job:
main CI excludes `arch_validation` and nightly L2 only scans
`tests/arch_validation_test.py`. It also downloaded the config from HuggingFace
(`AutoConfig.from_pretrained("deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True)`).
Make it a self-contained regression guard: build a tiny DeepSeek-V2-Lite-shaped
model from an inline `ArchitectureConfig` (MLA attention, first_k_dense_replace,
softmax routing, shared experts) with no network access, and drop the
`arch_validation` marker so it runs in the offline main CI. hidden_size and
moe_intermediate_size are multiples of the quantization group_size (128) as QMoE
requires; small MLA/dense dims are fine for MatMulNBits (ceil-padded K).
The invariant is unchanged: one com.microsoft::QMoE per routed MoE layer, shared
experts stay dense MatMulNBits, QMoE float routing inputs (1/3/6/14), and no
`.moe.experts.` initializers leak into the graph.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
src/mobius/_configs/_base.py:312
- Same as above:
# ruff:ignore[unused-import]is likely not honored by Ruff. Prefer# noqa: F401for this intentional side-effect import to avoid lint failures.
from mobius._configs import per_model # ruff:ignore[unused-import] - side effect
from mobius._configs._extractors import extract_audio_config as _dispatch
| registered with :mod:`mobius._configs._extractors` at import time. | ||
| """ | ||
| from mobius._configs import per_model # noqa: F401 - side-effect import | ||
| from mobius._configs import per_model # ruff:ignore[unused-import] - side effect |
There was a problem hiding this comment.
Fixed in a387ee2. Replaced the non-standard `# ruff:ignore[unused-import]` on both lines (298 and 311) with the repo-standard `# noqa: F401` (with a short "imported for registration side effect" note). This clears the RUF102 (invalid-rule-code) and F401 (unused-import) warnings that were failing Lint. lintrunner --all-files is now clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (3)
src/mobius/_configs/_base.py:298
# ruff:ignore[unused-import]is not a recognized Ruff suppression in this repo (and there are no other# ruff:directives elsewhere), so this likely won’t silence the intended unused-import warning for the side-effect import.
from mobius._configs import per_model # ruff:ignore[unused-import] - side effect
src/mobius/_configs/_base.py:311
# ruff:ignore[unused-import]is not a recognized Ruff suppression in this repo (and there are no other# ruff:directives elsewhere), so this likely won’t silence the intended unused-import warning for the side-effect import.
from mobius._configs import per_model # ruff:ignore[unused-import] - side effect
src/mobius/components/_moe.py:254
- MoELayer’s QMoE fastpath is enabled for any gate that implements
qmoe_routing. However,TopKGate.qmoe_routing’s docstring explicitly calls out that CUDA QMoE ignoresrouter_weightsand always uses softmax-top-k onrouter_probs, which cannot reproduceSigmoidTopKGate.forward(sigmoid-first routing). To avoid silently changing routing behavior on CUDA, SigmoidTopKGate should not go through the QMoE path.
if self._qmoe_quantization is not None and hasattr(self.gate, "qmoe_routing"):
self.experts = None
self._init_qmoe_parameters(expert_config)
| weight_t = op.Transpose(self.weight, perm=[1, 0]) | ||
| router_logits = op.MatMul(hidden_states, weight_t) | ||
| router_logits = op.Cast(router_logits, to=ir.DataType.FLOAT.value) | ||
| routing_probs = op.Softmax(router_logits, axis=-1) | ||
| return routing_probs, routing_probs, self.norm_topk_prob, 1.0 |
There was a problem hiding this comment.
Fixed in a387ee2. SoftmaxTopKGate.qmoe_routing (and SigmoidTopKGate.qmoe_routing on the other flagged line) now pass the raw float32 logits as router_probs (QMoE input 1) and the activated probabilities as router_weights (input 14), instead of passing the activated probs as both.
This matches the EP semantics established for TopKGate: CUDA QMoE ignores router_weights and runs softmax-top-k on router_probs, so feeding pre-activated probs there was a softmax-of-softmax (softmax-of-sigmoid for the sigmoid gate). Feeding logits reproduces forward on both EPs:
- CUDA:
softmax(logits)→ top-k → renormalize ==SoftmaxTopKGate.forward. - CPU: top-k selection over logits is identical (softmax/sigmoid are monotonic), and
router_weights(the activated probs) is gathered at the selected experts and renormalized ==forward.
The DeepSeek-V2-Lite int4 QMoE regression test (_group_query_attention_test.py, softmax routing) still passes — its FLOAT dtype/presence assertions on inputs 1/14 are unaffected. pytest src/mobius/components/_moe_test.py src/mobius/rewrite_rules/_group_query_attention_test.py -q → green.
- Replace invalid '# ruff:ignore[unused-import]' directives in _configs/_base.py with the repo-standard '# noqa: F401' for the per_model side-effect imports (fixes RUF102 + F401 Lint failures). - Apply ruff-format to PR-touched files (_moe_test, _weight_utils_test, deepseek, _moe) that had line-length reflow violations. - SoftmaxTopKGate/SigmoidTopKGate.qmoe_routing now pass raw float32 logits as the QMoE router_probs (input 1) and the activated probabilities as router_weights (input 14). CUDA QMoE ignores router_weights and applies softmax-top-k on router_probs, so passing activated probs there double-activated (softmax-of- softmax / softmax-of-sigmoid); logits reproduce forward on both EPs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Pushed a387ee2 to get Lint and test green: Lint — now clean ( Unit test — the only remaining failures in the Review comments addressed in code + replied on-thread: the Note: overlaps with #404 (GLM-5.2 MoE export), which also edits |
| assert config.num_local_experts is not None | ||
| assert config.num_experts_per_tok is not None | ||
| self.num_experts = config.num_local_experts | ||
| self.top_k = config.num_experts_per_tok | ||
| self._qmoe_quantization = _supported_qmoe_quantization(config.quantization) |
There was a problem hiding this comment.
Fixed in 98153fd. Updated the MoELayer class docstring to document both dispatch paths: the default loop-over-experts MLP path, and the fused com.microsoft::QMoE path (experts=None) selected when the quantization config matches the native QMoE ABI (_supported_qmoe_quantization) and the gate implements qmoe_routing. Also notes that the QMoE path packs quantized fc1/fc2 parameters instead of per-expert MLP modules, so the weight/preprocess expectations are clear.
…dispatch MoELayer now switches to a fused com.microsoft::QMoE path (experts=None) when the quantization config matches the native QMoE ABI and the gate implements qmoe_routing, otherwise it uses the loop-over-experts MLP dispatch. Update the class docstring to describe both paths so the weight/preprocess expectations are clear when debugging. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
| def qmoe_routing(self, op: OpBuilder, hidden_states: ir.Value): | ||
| """Return distinct expert-selection and aggregation scores for QMoE.""" | ||
| scores, scores_for_choice = self._routing_scores(op, hidden_states) | ||
| return ( | ||
| scores_for_choice, | ||
| scores, | ||
| self.norm_topk_prob, | ||
| float(self.routed_scaling_factor), | ||
| ) |
Summary
transformer.encoderGPTQ checkpoint names to Mobius canonical initializersValidation
pytest src/mobius/models/chatglm_test.py src/mobius/_configs_test.py -q(83 passed)pytest tests/build_graph_test.py -k chatglm -q(4 passed)Artifact:
/home/justinchu/glm-e2e-artifacts/glm-4-9b-int4-cuda-nosplitTokenizer note: the source repository only provides a custom slow tiktoken tokenizer. The smoke package uses the compatible fast
tokenizer.jsonfromTHUDM/glm-4-9b-chat-hf.This PR is intentionally left draft for Justin to review and merge.