Fix real PLaMo2 cached decode parity - #620
Merged
Merged
Conversation
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
justinchuby
force-pushed
the
justinchuby-validate-real-plamo2
branch
from
August 25, 2026 17:39
1315932 to
52fa76a
Compare
Honor PLaMo2's published local RoPE base when importing the legacy public GGUF metadata profile, and preserve A_log through runtime discretization. Add tied-weight handling, generic ORT GenAI decoder metadata, and real L4/L5 regression evidence for exact cached generation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Align PLaMo2's public recurrent-state ABI with ORT GenAI 0.15.2 while preserving sparse alternating layer indices. Package the immutable custom tokenizer sources, propagate remote-code trust during chat-template discovery, and record exact token-ID generation plus the tokenizer format limitation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchuby-fix-plamo2-runtime-parity
branch
from
August 25, 2026 17:45
e4184a0 to
0de76c6
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes real-model cached-decode parity for PLaMo2 by correcting GGUF/HF metadata mismatches (local RoPE base and Mamba decay handling), aligning the exported state ABI with released ORT GenAI behavior, and committing durable L4/L5 evidence and goldens for regression coverage.
Changes:
- Restore PLaMo2’s local RoPE base to
1e4for legacy GGUF metadata that incorrectly declares1e6, with a targeted GGUF postprocess override + regression test. - Preserve
A_login the model and compute-exp(A_log)in-graph, updating weight preprocessing and parity tests to match HF/GGUF discretization behavior. - Rename PLaMo2’s public SSM state I/O to
recurrent_stateand update ORT GenAI export/config handling plus testdata (goldens, evidence, L4/L5 case).
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/build_graph_test.py | Adjusts cache-output assertions so PLaMo2 mamba layers expect recurrent_state instead of ssm_state. |
| testdata/golden/causal-lm/plamo2-1b.json | Adds committed prefill golden IDs/logit summaries for PLaMo2. |
| testdata/golden/causal-lm/plamo2-1b_generation.json | Adds committed 20-token greedy generation golden for PLaMo2. |
| testdata/evidence/causal-lm/plamo2-1b-real.json | Updates durable real-model evidence: parity metrics, root-cause/fix narrative, and ORT GenAI ABI findings. |
| testdata/cases/causal-lm/plamo2-1b.yaml | Promotes/updates the L4+L5 case prompt and notes to reflect passing raw ORT + token-ID ORT GenAI results. |
| src/mobius/tasks/_plamo2.py | Renames exported SSM state I/O to recurrent_state and updates cache ABI/runtime support metadata. |
| src/mobius/models/plamo2.py | Switches decay storage to A_log, computes -exp(A_log) in-graph, and updates weight preprocessing for HF/GGUF + effective tying. |
| src/mobius/models/plamo2_test.py | Updates parity tests for A_log and recurrent_state, and adds tests for effective tied embedding behavior under quantization metadata. |
| src/mobius/integrations/ort_genai/auto_export.py | Maps plamo2 to generic ORT GenAI decoder, copies extra tokenizer artifacts, and propagates trust_remote_code to chat-template extraction. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Adds tests for the plamo2 model type mapping, trust-remote-code propagation, and copying of new tokenizer artifacts. |
| src/mobius/integrations/gguf/_config_mapping.py | Overrides rope_theta to 1e4 when legacy GGUF metadata uses 1e6, with explanatory comments. |
| src/mobius/integrations/gguf/_builder_test.py | Adds GGUF writer parameterization for RoPE base and verifies the legacy-million-base override; updates state name expectations to recurrent_state. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+431
to
438
| if name == "model.embed_tokens.weight" and tied_embeddings: | ||
| result[name] = value | ||
| # onnxscript materializes the shared Parameter under both use | ||
| # sites, while the official checkpoint stores only the embedding. | ||
| result["lm_head.weight"] = value | ||
| continue | ||
| if name == "lm_head.weight" and tied_embeddings: | ||
| continue |
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - clean up still-valid typing, documentation, test robustness, error-message, maintainability, and performance findings left by the GGUF PR stack - preserve the behavioral fixes merged in #625-#632 - hash reused multi-GB GGUF sources once, at the final pre-publication integrity gate, while retaining cheap identity checks around staging Exact base: `2db9d33debdc254d879a51b14434c9a81c230f4f` Exact head: `4541ed2bc9d2ab4227484510b4a85b2d9113eb25` ## Reconstructed original 17-item low-priority tranche The persisted audit retained only the totals, so this list was reconstructed from the live threads and current source. All but the already-fixed #596 comment are addressed in this PR. | PR | Comment | Disposition | |---|---:|---| | #550 | 3837144656 | Implemented: correct tuple return annotation | | #552 | 3837389183 | Implemented: multichannel waveform shape docs | | #559 | 3837716261 | Implemented: name-based cache assertions | | #573 | 3854308350 | Implemented: one final GGUF hash, with integrity regression coverage | | #574 | 3854345521 | Implemented: `TensorRole \| None` typing | | #574 | 3854345597 | Implemented: removed obsolete verdict filtering | | #577 | 3854472859 | Implemented: documented SSM sequence length | | #578 | 3843112613 | Implemented: documented F64 passthrough | | #579 | 3854518332 | Implemented: generalized fused-projection error | | #580 | 3854576300 | Implemented: removed brittle node counts | | #583 | 3854681386 | Implemented: documented conditional draft outputs | | #587 | 3854816872 | Implemented: corrected MTP output contract docs | | #596 | 3846110059 | Already fixed on base: unambiguous GQA bias comment | | #600 | 3855174281 | Implemented: metadata-count-only MTP error | | #607 | 3855776048 | Implemented: stable route-field assertions | | #607 | 3855776086 | Implemented: public tensor iterator | | #609 | 3856486136 | Implemented: fail-closed LM-head comment | ## Current unresolved-thread disposition This covers all 48 Copilot threads returned by the reproducible #600-#630 query. The one human #623 thread is excluded. | PR | Comment | Current-main disposition and evidence | |---|---:|---| | #600 | 3855174177 | Already fixed by #629: package cycle and reserved-sidecar validation | | #600 | 3855174238 | Already fixed by #629: explicit MTP sidecar naming/loading | | #600 | 3855174281 | Implemented here: error no longer invents an observed block count | | #602 | 3848155961 | Outside exact stack; already fixed: top-level `expert_dtype` is classified before early return | | #602 | 3848155983 | Outside exact stack; still-valid behavioral block-quant validation, unchanged | | #602 | 3848156000 | Outside exact stack; still-valid truncated-read behavioral finding, unchanged | | #602 | 3848156022 | Outside exact stack; still-valid descriptor byte/dtype validation, unchanged | | #602 | 3848156040 | Outside exact stack; still-valid expert-bank payload validation, unchanged | | #603 | 3855249765 | Already fixed by #630: runtime preflight preserves shard sets | | #603 | 3855249840 | Already fixed by #630: success output follows durable runtime publication | | #604 | 3855343082 | Implemented here: graph-only MTP persistence distinguished from runtime rejection | | #604 | 3855343131 | Already fixed by #630: runtime success messages are atomic | | #607 | 3855776001 | Implemented here: missing generation golden skips before provenance read | | #607 | 3855776048 | Implemented here: only stable route fields are asserted | | #607 | 3855776086 | Implemented here: tensor count uses `tensor_items_raw()` | | #608 | 3856079371 | Still-valid behavioral cache-symlink containment finding; unchanged | | #608 | 3856079415 | Still-valid behavioral lowercase-digest validation finding; unchanged | | #609 | 3856486136 | Implemented here: comment matches value-preserving policy | | #610 | 3855541683 | Already fixed by #628: Falcon bias precedence is explicit | | #610 | 3855541761 | Already fixed by #628: CTRL tiny config exercises projection biases | | #611 | 3856595840 | Implemented here: runtime test resolves the distribution providing the module | | #612 | 3855677383 | Already fixed by #625: supported-version endianness detection | | #612 | 3855677427 | Implemented here: shared `INT64_MAX` sentinel | | #612 | 3855677460 | Implemented here: shared PLaMo2 width inference | | #612 | 3855677486 | Implemented here: accepted PLaMo2 activation spellings are explicit | | #613 | 3855845678 | Implemented here: canonical issue URL | | #613 | 3855845757 | Implemented here: Mamba-1 function-registration docs | | #613 | 3855845806 | Implemented here: test expects the canonical issue URL | | #614 | 3855988545 | Implemented here: removed stale Nemotron-H divergence comments | | #614 | 3855988597 | Already fixed by #628: zero-head geometry raises actionable `ValueError` | | #615 | 3856082290 | Already fixed by #626: dense GraniteHybrid bias closure | | #618 | 3856729330 | Implemented here: required routes filter ORT GenAI evidence | | #618 | 3856729409 | Implemented here: env-selected runtime version is authoritative | | #618 | 3856729490 | Stale/N/A: PR-description-only matrix claim; repository workflow claims one pinned version | | #618 | 3856729563 | Implemented here: schema tail restored to normal indentation | | #619 | 3856342484 | Already fixed on base: Kimi Linear uses `/issues/605` | | #619 | 3856342532 | Already fixed by #628: config rejects convolution kernels below 2 | | #619 | 3856342580 | Already fixed by #628: GGUF contract rejects convolution kernels below 2 | | #620 | 3855717041 | Already fixed by #627: tied LM-head-only checkpoints are retained | | #621 | 3856722324 | Already fixed by #628: Kimi-K3 required metadata is complete | | #623 | 3855931210 | N/A to current main: comment belongs to open, unmerged #623 | | #623 | 3855939302 | N/A to current main: comment belongs to open, unmerged #623 | | #623 | 3855939358 | N/A to current main: comment belongs to open, unmerged #623 | | #623 | 3855939394 | N/A to current main: comment belongs to open, unmerged #623 | | #624 | 3856777152 | Implemented here: runtime compatibility reuses the emitted model type | | #625 | 3857049733 | Implemented here: unsupported header reports both endian candidates | | #629 | 3857313182 | Newer post-audit behavioral sidecar-symlink cleanup finding; unchanged | | #629 | 3857313251 | Newer post-audit cross-platform path-safety finding; unchanged | ## Validation - affected GGUF/package/ORT GenAI/model/schema tests: 1,040 passed - broad non-integration suite: 7,851 passed, 56 skipped, 1 subtest passed - generated GGUF docs checks: 7 passed - initialized `lintrunner`; full lint/format passed - GPT-5.6 Sol medium review: one integrity finding fixed; re-review found no significant issues Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
1e4local RoPE base for the legacy public GGUF profile whose metadata incorrectly declares1e6A_logand evaluate-exp(A_log)in the ONNX graph, including correct HF/GGUF and effectively-tied weight loadingRoot cause
The pinned GGUF declares
plamo2.rope.freq_base=1e6, while the pinned HF architecture hasfull_attention_idx=[]and usesrope_local_theta=1e4for every attention layer. Trusting the contradictory GGUF field caused the first divergence in layer 1 attention during cached decode; later recurrent layers amplified it. Mamba state update, convolution orientation, attention QKV/QK norm, residual ordering, cache indexing, fusion, and rollback/replay were not the source.Real-model validation
Pinned revisions:
pfnet/plamo-2-1b:92c75fd6eea9018bcb9c33ee8921589febe071fa184326afcef7232d94df09a33279598255890325Raw ORT 1.29 CPU after the fix:
5.34057617e-052.47955322e-054.62532043e-05ORT GenAI 0.15.2 source independently discovers sparse KV indices and sparse
conv_stateindices, then derives the pairedrecurrent_statenames. After renaming PLaMo2's semantic SSM state to that public ABI, a CPU-specialized package with all eight attention layers lowered to GQA loaded and generated the exact verified 20-token continuation from prompt IDs[1, 6721]. UnspecializedAttentionplus recurrent state is still rejected during package config generation.The immutable tokenizer source was audited against the GGUF: all 100,000 ordered tokens, token types, float32 scores, and special IDs match exactly, with no normalizer or pre-tokenizer. Text tokenization remains deferred because ORT GenAI requires
tokenizer.json, and standard Unigram byte fallback is not identity-preserving: the official tokenizer encodes literal<0xFA>as[60, 48, 120, 5109, 62], while the standard tokenizer emits byte token[250].Branch reconstruction
Parent #617 was admin squash-merged as
79efe2c4512a40900ac5af5e794d057ea39c8f5b. This branch was reconstructed directly on that commit and retains only the two reviewed #620 commits, preserving #617 as the explicit parent while removing its duplicated pre-squash history.Tests