From 846b298ed3eeb405824fd3885c2c28c84a97f347 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 25 Aug 2026 00:43:13 -0700 Subject: [PATCH 1/2] Fix real PLaMo2 cached decode parity 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 --- src/mobius/integrations/gguf/_builder_test.py | 13 +- .../integrations/gguf/_config_mapping.py | 8 ++ .../integrations/ort_genai/auto_export.py | 3 + .../ort_genai/auto_export_test.py | 3 + src/mobius/models/plamo2.py | 24 +++- src/mobius/models/plamo2_test.py | 41 ++++-- testdata/cases/causal-lm/plamo2-1b.yaml | 5 +- .../evidence/causal-lm/plamo2-1b-real.json | 135 +++++++++++------- testdata/golden/causal-lm/plamo2-1b.json | 38 +++++ .../causal-lm/plamo2-1b_generation.json | 27 ++++ 10 files changed, 225 insertions(+), 72 deletions(-) create mode 100644 testdata/golden/causal-lm/plamo2-1b.json create mode 100644 testdata/golden/causal-lm/plamo2-1b_generation.json diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index 2a5e77850..d5805f6ea 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -840,6 +840,7 @@ def _write_plamo2_gguf( legacy_scalar_heads: bool = False, quantized_embedding: bool = False, include_output: bool = False, + rope_theta: float = 10_000.0, ) -> None: """Write a complete tiny alternating PLaMo2 GGUF.""" from gguf import GGMLQuantizationType, GGUFWriter @@ -867,7 +868,7 @@ def _write_plamo2_gguf( kv_heads if legacy_scalar_heads else kv_head_counts or [0, kv_heads] ) writer.add_layer_norm_rms_eps(epsilon) - writer.add_rope_freq_base(10_000.0) + writer.add_rope_freq_base(rope_theta) writer.add_vocab_size(vocab) writer.add_ssm_conv_kernel(kernel) writer.add_ssm_inner_size(inner) @@ -4560,6 +4561,14 @@ def test_legacy_scalar_heads_infer_exact_tensor_schedule(self, tmp_path: Path) - assert package.config.attention_head_counts == (0, 4) assert package.config.attention_kv_head_counts == (0, 2) + def test_legacy_million_base_restores_reference_local_rope(self, tmp_path: Path) -> None: + from mobius.integrations.gguf import build_from_gguf + + path = tmp_path / "plamo2-legacy-rope.gguf" + _write_plamo2_gguf(path, quantized=False, rope_theta=1_000_000.0) + package = build_from_gguf(path) + assert package.config.rope_theta == pytest.approx(10_000.0) + def test_quantized_source_preserves_exact_projection_roles(self, tmp_path: Path) -> None: from mobius.integrations.gguf import build_from_gguf @@ -4568,7 +4577,7 @@ def test_quantized_source_preserves_exact_projection_roles(self, tmp_path: Path) model = build_from_gguf(path, keep_quantized=True)["model"] assert sum(node.op_type == "MatMulNBits" for node in model.graph) == 9 assert ( - model.graph.initializers["model.layers.0.mixer.A"].const_value.dtype + model.graph.initializers["model.layers.0.mixer.A_log"].const_value.dtype == ir.DataType.FLOAT ) assert ( diff --git a/src/mobius/integrations/gguf/_config_mapping.py b/src/mobius/integrations/gguf/_config_mapping.py index afa4dd0d9..2dc8e7580 100644 --- a/src/mobius/integrations/gguf/_config_mapping.py +++ b/src/mobius/integrations/gguf/_config_mapping.py @@ -2244,6 +2244,14 @@ def _plamo2_postprocess( fields.update( hidden_act="silu", head_dim=key_length, + # The released PLaMo2 converter wrote 1e6 here even though the pinned + # reference architecture uses its 1e4 local-RoPE default for every + # attention layer. Preserve other explicit bases for future variants. + rope_theta=( + 10_000.0 + if float(metadata[f"{arch}.rope.freq_base"]) == 1_000_000.0 # noqa: RUF069 + else config.rope_theta + ), attention_head_counts=head_counts, attention_kv_head_counts=kv_head_counts, mamba_num_heads=ssm_heads, diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 926229fb8..ad7270483 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -108,6 +108,9 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]: # ORT GenAI (see onnxruntime-genai/src/models/model_type.h LLM list). "hunyuan_v1_dense": "decoder", "deepseek_v4": "decoder", + # PLaMo2 is a decoder-only hybrid. Released ORT GenAI does not have a + # model-specific registry entry, so emit its generic decoder type. + "plamo2": "decoder", # Qwen VL model families have separate ORT GenAI model types. "qwen2_vl": "qwen2_5_vl", "qwen3_vl": "qwen3_vl", diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 0e3263663..d75bca280 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -110,6 +110,9 @@ def test_hunyuan_v1_dense_maps_to_decoder(self): # decoder-only causal LM not in its built-in registry. assert _resolve_ort_genai_model_type("hunyuan_v1_dense") == "decoder" + def test_plamo2_maps_to_generic_decoder(self): + assert _resolve_ort_genai_model_type("plamo2") == "decoder" + def test_unknown_model_type_passthrough(self): assert _resolve_ort_genai_model_type("my_custom") == "my_custom" diff --git a/src/mobius/models/plamo2.py b/src/mobius/models/plamo2.py index 8149e4c12..00ec299c8 100644 --- a/src/mobius/models/plamo2.py +++ b/src/mobius/models/plamo2.py @@ -140,8 +140,7 @@ def __init__(self, config: Plamo2Config): # dt participates in Softplus and remains float even for quantized GGUF imports. self.dt_proj = Linear(self.dt_rank, self.num_heads, bias=False) self.dt_bias = nn.Parameter([self.num_heads]) - # GGUF stores the already transformed negative decay directly. - self.A = nn.Parameter([self.num_heads]) + self.A_log = nn.Parameter([self.num_heads]) self.D = nn.Parameter([self.num_heads]) self.dt_norm_weight = nn.Parameter([self.dt_rank]) self.B_norm_weight = nn.Parameter([self.state_size]) @@ -198,7 +197,10 @@ def forward( ) ) dt = op.Mul(dt, op.CastLike(padding_mask, dt)) - decay = op.Mul(dt, op.Cast(self.A, to=ir.DataType.FLOAT)) + # Match the reference's float32 discretization at runtime rather than + # baking a rounded -exp(A_log) into imported weights. + a = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT))) + decay = op.Mul(dt, a) x_f32 = op.Cast(x, to=ir.DataType.FLOAT) x_heads = op.Reshape(x_f32, [0, 0, self.num_heads, self.head_dim]) value = op.Reshape( @@ -414,7 +416,7 @@ def preprocess_weights( self, state_dict: dict[str, torch.Tensor], ) -> dict[str, torch.Tensor]: - """Convert the official offset norms and Mamba decay to graph values.""" + """Convert official and GGUF offset norms and Mamba decay names.""" result: dict[str, torch.Tensor] = {} norm_offsets = { ".pre_mixer_norm.weight": 1.0, @@ -423,9 +425,16 @@ def preprocess_weights( ".post_mlp_norm.weight": 1.0 / (5.0**1.5), } norms_are_folded = bool(getattr(self.config, "_plamo2_norms_are_folded", False)) + tied_embeddings = effective_tie_word_embeddings(self.config) for name, value in state_dict.items(): name = name.replace("model.layers.layers.", "model.layers.") - if name == "lm_head.weight" and self.config.tie_word_embeddings: + 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 if name == "model.norm.weight": result[name] = value if norms_are_folded else value + 1.0 @@ -436,8 +445,9 @@ def preprocess_weights( ) if offset is not None: result[name] = value if norms_are_folded else value + offset - elif name.endswith(".mixer.A_log"): - result[name.removesuffix("A_log") + "A"] = -torch.exp(value) + elif name.endswith(".mixer.A"): + # llama.cpp serializes PLaMo2's A_log as -exp(A_log). + result[name.removesuffix("A") + "A_log"] = torch.log(-value) else: result[name] = value return result diff --git a/src/mobius/models/plamo2_test.py b/src/mobius/models/plamo2_test.py index 904fae5f3..507b76ff5 100644 --- a/src/mobius/models/plamo2_test.py +++ b/src/mobius/models/plamo2_test.py @@ -10,7 +10,7 @@ import torch from mobius import build_from_module -from mobius._configs import Plamo2Config +from mobius._configs import Plamo2Config, QuantizationConfig from mobius._testing.ort_inference import OnnxModelSession from mobius.models.plamo2 import Plamo2ForCausalLM from mobius.tasks import Plamo2CausalLMTask @@ -50,9 +50,7 @@ def _fill_weights(model: ir.Model) -> None: continue shape = [dim if isinstance(dim, int) else 1 for dim in initializer.shape] values = (rng.standard_normal(shape) * 0.03).astype(initializer.dtype.numpy()) - if initializer.name.endswith(".A"): - values = -np.exp(values) - elif "norm" in initializer.name: + if "norm" in initializer.name: values += 1.0 initializer.const_value = ir.Tensor(values) @@ -124,10 +122,8 @@ def test_plamo2_mamba_matches_independent_reduced_reference() -> None: initializer.const_value = ir.Tensor(value) values[initializer.name] = value - def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndarray: + def assign(name: str, scale: float = 0.05) -> np.ndarray: value = (rng.standard_normal(values[name].shape) * scale).astype(np.float32) - if negative: - value = -np.exp(value) model.graph.initializers[name].const_value = ir.Tensor(value) values[name] = value return value @@ -138,7 +134,7 @@ def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndar assign("model.layers.0.mixer.bcdt_proj.weight") assign("model.layers.0.mixer.dt_proj.weight") assign("model.layers.0.mixer.dt_bias") - assign("model.layers.0.mixer.A", negative=True) + assign("model.layers.0.mixer.A_log") assign("model.layers.0.mixer.D") assign("model.layers.0.mixer.out_proj.weight") @@ -204,7 +200,8 @@ def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndar scan_outputs = [] for index in range(input_ids.shape[1]): decay = np.exp( - dt[:, index, :, None, None] * values["model.layers.0.mixer.A"][None, :, None, None] + dt[:, index, :, None, None] + * -np.exp(values["model.layers.0.mixer.A_log"])[None, :, None, None] ) update = ( dt[:, index, :, None, None] @@ -451,6 +448,7 @@ def test_plamo2_preprocesses_offsets_and_decay_values() -> None: "model.layers.layers.0.pre_mlp_norm.weight": torch.tensor([5.0]), "model.layers.layers.0.post_mlp_norm.weight": torch.tensor([6.0]), "model.layers.layers.0.mixer.A_log": torch.tensor([0.0, 1.0]), + "model.embed_tokens.weight": torch.tensor([8.0]), "lm_head.weight": torch.tensor([9.0]), } actual = model.preprocess_weights(weights) @@ -468,10 +466,31 @@ def test_plamo2_preprocesses_offsets_and_decay_values() -> None: actual["model.layers.0.post_mlp_norm.weight"], torch.tensor([6.0 + 1.0 / (5.0**1.5)]), ) + torch.testing.assert_close(actual["model.layers.0.mixer.A_log"], torch.tensor([0.0, 1.0])) + gguf_actual = model.preprocess_weights( + {"model.layers.0.mixer.A": -torch.exp(torch.tensor([0.0, 1.0]))} + ) torch.testing.assert_close( - actual["model.layers.0.mixer.A"], -torch.exp(torch.tensor([0.0, 1.0])) + gguf_actual["model.layers.0.mixer.A_log"], torch.tensor([0.0, 1.0]) ) - assert "lm_head.weight" not in actual + torch.testing.assert_close(actual["lm_head.weight"], torch.tensor([8.0])) + + +def test_plamo2_preprocesses_effectively_tied_quantized_embeddings() -> None: + config = _config() + config.tie_word_embeddings = False + config.quantization = QuantizationConfig(tie_word_embeddings=True) + model = Plamo2ForCausalLM(config) + + actual = model.preprocess_weights( + { + "model.embed_tokens.weight": torch.tensor([8.0]), + "lm_head.weight": torch.tensor([9.0]), + } + ) + + torch.testing.assert_close(actual["model.embed_tokens.weight"], torch.tensor([8.0])) + torch.testing.assert_close(actual["lm_head.weight"], torch.tensor([8.0])) def test_plamo2_transformers_config_uses_explicit_schedule_and_pinned_defaults() -> None: diff --git a/testdata/cases/causal-lm/plamo2-1b.yaml b/testdata/cases/causal-lm/plamo2-1b.yaml index 7cadef64b..cfd0316d1 100644 --- a/testdata/cases/causal-lm/plamo2-1b.yaml +++ b/testdata/cases/causal-lm/plamo2-1b.yaml @@ -7,7 +7,7 @@ trust_remote_code: true inputs: prompts: - - "The future of local inference is" + - "Hello" level: "L4+L5" @@ -15,5 +15,4 @@ generation: max_new_tokens: 20 do_sample: false -notes: "The exact public checkpoint/GGUF route was exercised; see testdata/evidence/causal-lm/plamo2-1b-real.json." -skip_reason: "Real CPU validation found cached-decode divergence, and ORT GenAI 0.15.2 cannot bind the heterogeneous recurrent/KV state ABI; do not claim L4/L5 yet." +notes: "The pinned source checkpoint and public F32 GGUF pass raw ORT prefill, all cached decode steps, rollback/replay, and exact 20-token generation. ORT GenAI remains deferred because 0.15.2 cannot bind the heterogeneous recurrent/KV state ABI." diff --git a/testdata/evidence/causal-lm/plamo2-1b-real.json b/testdata/evidence/causal-lm/plamo2-1b-real.json index ee43300c0..75cf4914b 100644 --- a/testdata/evidence/causal-lm/plamo2-1b-real.json +++ b/testdata/evidence/causal-lm/plamo2-1b-real.json @@ -1,9 +1,9 @@ { "schema_version": 1, - "validated_at_utc": "2026-08-25T06:14:35Z", - "validated_parent_commit": "64aa8fb55d10d3097d944e44231f0301d2527113", + "validated_at_utc": "2026-08-25T07:43:02Z", + "validated_parent_commit": "1315932c93753931aa271d460cc9d01af6a3f9d1", "validation_change_set": "the compatibility and evidence changes committed with this record", - "status": "runtime-deferred", + "status": "raw-runtime-passed-ort-genai-deferred", "source_checkpoint": { "repository": "pfnet/plamo-2-1b", "revision": "92c75fd6eea9018bcb9c33ee8921589febe071fa", @@ -65,12 +65,12 @@ "provenance_comparison": { "method": "streamed comparison of every GGUF tensor against the pinned safetensors after documented norm-offset, convolution-shape, and Mamba-decay transforms", "compared_tensors": 218, - "exact_tensors": 213, + "exact_tensors": 211, "missing_tensors": 0, "elements": 1291441920, - "max_abs_error": 9.5367431640625e-07, - "mean_abs_error": 3.323056466858475e-15, - "non_exact_tensors": "five ssm_a tensors differ only by float32 -exp(A_log) rounding" + "max_abs_error": 0.0001220703125, + "mean_abs_error": 2.4659848197812267e-13, + "non_exact_tensors": "seven ssm_a tensors differ after llama.cpp eagerly applies float32 -exp(A_log); the largest absolute difference is one ULP at |A| ~= 1024" }, "architecture": { "layers": 16, @@ -100,10 +100,11 @@ "transformers_version": "5.12.1", "torch_version": "2.12.1", "compatibility_shims": [ - "adapt Transformers 5 tied-weight metadata and tie lm_head to embed_tokens", + "adapt Transformers 5 tied-weight metadata, assign every shard on meta, verify all 218 parameters exactly, then tie lm_head to embed_tokens", "rebuild non-persistent RoPE caches after meta-device loading", "provide pure-PyTorch causal convolution and selective state update" ], + "superseded_result": "the previous reference silently retained initialized parameters under Transformers 5.12.1; its plausible generated text was not checkpoint output", "prompt": "Hello", "prompt_tokens": [ 1, @@ -111,39 +112,38 @@ ], "generated_tokens": [ 44, - 24514, - 45119, - 7982, + 2268, + 1101, + 1080, + 1087, + 8789, 45114, - 2166, - 353, - 1468, - 15658, - 1090, - 2596, + 2002, + 3104, + 1131, + 1385, 1031, 73, - 1170, - 1090, - 2596, - 45114, - 2166, + 39, + 109, + 2569, + 3648, 353, - 1468 + 32, + 79135 ], - "generated_text": ", my name is Alex and I'm a software engineer at Google.\n\nI work at Google and I'm a software", + "generated_text": ", I'm new to this forum and am looking for some help.\n\nI'm trying to install a 'm trying to Raspberry Pi", "generated_length": 20, "rollback_replay_exact": true }, - "ort_cpu": { + "source_ort_cpu": { "onnxruntime_version": "1.29.0", "execution_provider": "CPUExecutionProvider", - "fused_and_portable_results_equal": true, - "prefill_max_abs_error": 0.011388778686523438, - "prefill_mean_abs_error": 0.000737154199364013, + "prefill_max_abs_error": 0.0000534057617, + "prefill_mean_abs_error": 0.0000069309889, "decode_steps": 19, - "decode_max_abs_error": 23.476438999176025, - "decode_mean_abs_error": 1.668023172595157, + "decode_max_abs_error": 0.000057220459, + "decode_max_mean_abs_error": 0.00000550644416, "generated_tokens": [ 44, 2268, @@ -153,37 +153,74 @@ 8789, 45114, 2002, + 3104, + 1131, + 1385, + 1031, + 73, + 39, + 109, 2569, - 3035, - 8435, - 46, - 7525, - 2004, - 1078, - 1652, - 290, - 15143, - 45116, - 2057 + 3648, + 353, + 32, + 79135 ], "generated_length": 20, - "token_equality": false, + "token_equality": true, "rollback_replay_exact": true, - "promotion_decision": "blocked: cached decode does not meet numerical or token-equality gates" + "promotion_decision": "passed: prefill, every teacher-forced cached decode, and exact 20-token greedy generation match the verified checkpoint" + }, + "gguf_ort_cpu": { + "onnxruntime_version": "1.29.0", + "execution_provider": "CPUExecutionProvider", + "fused_and_portable_results_equal": true, + "prefill_max_abs_error": 0.0000534057617, + "prefill_mean_abs_error": 0.00000696335046, + "decode_steps": 19, + "decode_max_abs_error": 0.0000462532043, + "decode_max_mean_abs_error": 0.00000590607647, + "generated_tokens": [ + 44, + 2268, + 1101, + 1080, + 1087, + 8789, + 45114, + 2002, + 3104, + 1131, + 1385, + 1031, + 73, + 39, + 109, + 2569, + 3648, + 353, + 32, + 79135 + ], + "token_equality": true, + "root_cause": "the pinned GGUF declares plamo2.rope.freq_base=1e6, but the pinned reference has full_attention_idx=[] and uses its rope_local_theta=1e4 default for every attention layer; trusting the contradictory GGUF field caused the first cached divergence in layer 1 attention and amplified through later recurrent states", + "fix": "restore the published PLaMo2 local-RoPE base of 1e4 for the legacy 1e6 GGUF metadata profile and keep A_log in the graph so source checkpoints retain reference float32 discretization", + "promotion_decision": "passed: prefill, every teacher-forced cached decode, rollback/replay, and exact 20-token greedy generation match" }, "ort_genai": { "released_version": "0.15.2", - "generated_model_type": "plamo2", + "generated_model_type": "decoder", "generated_state_schema": "homogeneous past_key_names/past_value_names templates", - "model_load": "failed: Unsupported model_type in config.json: plamo2", - "tokenizer_load_after_forcing_llama": "failed: Invalid file: tokenizer.json", - "generation_after_forcing_llama_and_supplying_prompt_tokens": "failed: Model input was not found: past_key_values.0.recurrent_state", + "model_load": "passed with type=decoder and num_hidden_layers=16 discovered from sparse global state indices", + "tokenizer_load": "failed: the exact packaged tokenizer.jsonl format is not accepted by ORT GenAI 0.15.2", + "generation_with_prompt_tokens": "failed: Model input was not found: past_key_values.0.recurrent_state", "upstream_main_commit": "a8e0fdf81b061e67c1c3f9485bfdc06735ccd473", "upstream_main_install": "not feasible from the Python source package: pip built UNKNOWN-0.0.0 without an onnxruntime_genai module", "promotion_decision": "blocked: no released heterogeneous-state schema, no compatible tokenizer artifact, and no successful end-to-end generation" }, "l4_l5": { - "goldens_committed": false, - "reason": "real prefill exceeds the parity gate and cached generation diverges; committing passing goldens would misrepresent support" + "raw_ort_parity": "passed", + "generated_tokens_committed": true, + "ort_genai_e2e": "deferred because released ORT GenAI cannot represent the alternating state ABI" } } diff --git a/testdata/golden/causal-lm/plamo2-1b.json b/testdata/golden/causal-lm/plamo2-1b.json new file mode 100644 index 000000000..ec924ed19 --- /dev/null +++ b/testdata/golden/causal-lm/plamo2-1b.json @@ -0,0 +1,38 @@ +{ + "top1_id": 44, + "top2_id": 2151, + "top10_ids": [ + 44, + 2151, + 1102, + 45114, + 1029, + 1214, + 1462, + 33, + 47126, + 1422 + ], + "top10_logits": [ + "0x1.e1d81a0000000p+2", + "0x1.b9c1700000000p+2", + "0x1.9f871c0000000p+2", + "0x1.97007e0000000p+2", + "0x1.87430e0000000p+2", + "0x1.864ac00000000p+2", + "0x1.81f87e0000000p+2", + "0x1.78e0480000000p+2", + "0x1.6e2e900000000p+2", + "0x1.68f2820000000p+2" + ], + "logits_summary": [ + "0x1.e1d81a0000000p+2", + "-0x1.144fac0000000p+5", + "-0x1.dfeaae0000000p+2", + "0x1.4e42cc0000000p+1" + ], + "input_ids": [ + 1, + 6721 + ] +} diff --git a/testdata/golden/causal-lm/plamo2-1b_generation.json b/testdata/golden/causal-lm/plamo2-1b_generation.json new file mode 100644 index 000000000..0fc288d9f --- /dev/null +++ b/testdata/golden/causal-lm/plamo2-1b_generation.json @@ -0,0 +1,27 @@ +{ + "model_id": "pfnet/plamo-2-1b", + "prompt": "Hello", + "generated_tokens": [ + 44, + 2268, + 1101, + 1080, + 1087, + 8789, + 45114, + 2002, + 3104, + 1131, + 1385, + 1031, + 73, + 39, + 109, + 2569, + 3648, + 353, + 32, + 79135 + ], + "generated_text": ", I'm new to this forum and am looking for some help.\n\nI'm trying to install a 'm trying to Raspberry Pi" +} From 0de76c6a42efca8d48e624066a6c67a2c2bd6058 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 25 Aug 2026 01:07:41 -0700 Subject: [PATCH 2/2] Enable PLaMo2 ORT GenAI recurrent state 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 --- src/mobius/integrations/gguf/_builder_test.py | 10 ++--- .../integrations/ort_genai/auto_export.py | 11 ++++- .../ort_genai/auto_export_test.py | 18 ++++++++- src/mobius/models/plamo2_test.py | 10 +++-- src/mobius/tasks/_plamo2.py | 9 ++--- testdata/cases/causal-lm/plamo2-1b.yaml | 2 +- .../evidence/causal-lm/plamo2-1b-real.json | 40 ++++++++++++++++--- tests/build_graph_test.py | 5 ++- 8 files changed, 80 insertions(+), 25 deletions(-) diff --git a/src/mobius/integrations/gguf/_builder_test.py b/src/mobius/integrations/gguf/_builder_test.py index d5805f6ea..c565aa8fb 100644 --- a/src/mobius/integrations/gguf/_builder_test.py +++ b/src/mobius/integrations/gguf/_builder_test.py @@ -4505,7 +4505,7 @@ def _inputs(batch: int) -> dict[str, np.ndarray]: "position_ids": np.asarray([[0, 1], [0, 1]][:batch], np.int64), "attention_mask": np.ones((batch, 2), np.int64), "past_key_values.0.conv_state": np.zeros((batch, 32, 3), np.float32), - "past_key_values.0.ssm_state": np.zeros((batch, 4, 8, 4), np.float32), + "past_key_values.0.recurrent_state": np.zeros((batch, 4, 8, 4), np.float32), "past_key_values.1.key": np.zeros((batch, 2, 0, 8), np.float32), "past_key_values.1.value": np.zeros((batch, 2, 0, 8), np.float32), } @@ -4534,13 +4534,13 @@ def test_float_import_executes_and_round_trips(self, tmp_path: Path) -> None: model.graph.initializers[initializer_name].const_value.numpy(), source_tensors[source_name], ) - assert model.metadata_props["mobius.runtime_support"].endswith( - "onnxruntime/mobius#605" + assert model.metadata_props["mobius.runtime_support"] == ( + "ORT GenAI 0.15.2 state ABI; package requires GQA-specialized attention" ) assert [value.name for value in model.graph.outputs] == [ "logits", "present.0.conv_state", - "present.0.ssm_state", + "present.0.recurrent_state", "present.1.key", "present.1.value", ] @@ -4549,7 +4549,7 @@ def test_float_import_executes_and_round_trips(self, tmp_path: Path) -> None: session = OnnxModelSession(ModelPackage.load(output_dir)["model"]) outputs = session.run(self._inputs(2)) assert outputs["logits"].shape == (2, 2, 64) - assert outputs["present.0.ssm_state"].dtype == np.float32 + assert outputs["present.0.recurrent_state"].dtype == np.float32 def test_legacy_scalar_heads_infer_exact_tensor_schedule(self, tmp_path: Path) -> None: from mobius.integrations.gguf import build_from_gguf diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index ad7270483..9e1626612 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -205,6 +205,7 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]: _TOKENIZER_FILES = [ "tokenizer.json", + "tokenizer.jsonl", # PLaMo2 scored vocabulary "tokenizer_config.json", "special_tokens_map.json", "tokenizer.model", # SentencePiece @@ -212,6 +213,7 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]: "merges.txt", # BPE "vocab.json", # BPE "chat_template.jinja", # Chat template for ORT GenAI + "tokenization_plamo.py", # PLaMo2's exact custom tokenizer implementation # Preserve HuggingFace processor metadata for VLMs whose preprocessing # cannot be represented by an ort-extensions image_processor.json. "preprocessor_config.json", @@ -480,6 +482,7 @@ def _fix_chat_template( hf_model_id: str | None, *, revision: str | None = None, + trust_remote_code: bool = False, ) -> bool: """Ensure chat_template is present in tokenizer_config.json. @@ -511,6 +514,7 @@ def _fix_chat_template( tokenizer = AutoTokenizer.from_pretrained( hf_model_id, + trust_remote_code=trust_remote_code, **_revision_kwargs(revision), ) template = getattr(tokenizer, "chat_template", None) @@ -1763,7 +1767,12 @@ def write_ort_genai_config( _fix_tokenizer_config(directory) # Ensure chat_template is in tokenizer_config.json - _fix_chat_template(directory, hf_model_id, revision=revision) + _fix_chat_template( + directory, + hf_model_id, + revision=revision, + trust_remote_code=trust_remote_code, + ) # Correct assets that ship broken from upstream apply_asset_patches(directory) diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index d75bca280..65416e7a9 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -810,10 +810,20 @@ def test_adds_chat_template(self, tmp_path): with mock.patch( "transformers.AutoTokenizer.from_pretrained", return_value=fake_tokenizer, - ): - result = _fix_chat_template(str(tmp_path), "fake/model") + ) as from_pretrained: + result = _fix_chat_template( + str(tmp_path), + "fake/model", + revision="immutable-revision", + trust_remote_code=True, + ) assert result is True + from_pretrained.assert_called_once_with( + "fake/model", + revision="immutable-revision", + trust_remote_code=True, + ) fixed = json.loads((tmp_path / "tokenizer_config.json").read_text()) assert fixed["chat_template"] == "{{ bos_token }}" @@ -949,7 +959,9 @@ def test_copies_present_files(self, tmp_path): src = tmp_path / "model" src.mkdir() (src / "tokenizer.json").write_text('{"test": true}') + (src / "tokenizer.jsonl").write_text('["token", 0.0, "NORMAL"]\n') (src / "tokenizer_config.json").write_text('{"model_type": "llama"}') + (src / "tokenization_plamo.py").write_text("class Plamo2Tokenizer: pass\n") (src / "chat_template.jinja").write_text("{{ messages }}") dst = tmp_path / "output" @@ -958,7 +970,9 @@ def test_copies_present_files(self, tmp_path): assert set(copied) == { "tokenizer.json", + "tokenizer.jsonl", "tokenizer_config.json", + "tokenization_plamo.py", "chat_template.jinja", } assert (dst / "tokenizer.json").read_text() == '{"test": true}' diff --git a/src/mobius/models/plamo2_test.py b/src/mobius/models/plamo2_test.py index 507b76ff5..a5aaa1dc2 100644 --- a/src/mobius/models/plamo2_test.py +++ b/src/mobius/models/plamo2_test.py @@ -60,7 +60,7 @@ def _empty_states(config: Plamo2Config, batch: int) -> dict[str, np.ndarray]: "past_key_values.0.conv_state": np.zeros( (batch, config.mamba_inner_size, config.mamba_d_conv - 1), np.float32 ), - "past_key_values.0.ssm_state": np.zeros( + "past_key_values.0.recurrent_state": np.zeros( ( batch, config.mamba_num_heads, @@ -226,7 +226,9 @@ def assign(name: str, scale: float = 0.05) -> np.ndarray: rtol=1e-6, atol=1e-7, ) - np.testing.assert_allclose(outputs["present.0.ssm_state"], state, rtol=2e-5, atol=2e-6) + np.testing.assert_allclose( + outputs["present.0.recurrent_state"], state, rtol=2e-5, atol=2e-6 + ) def test_plamo2_attention_and_mlp_match_independent_reduced_reference() -> None: @@ -404,7 +406,9 @@ def test_plamo2_left_padding_does_not_change_recurrent_state_or_valid_logits() - padded["present.0.conv_state"], unpadded["present.0.conv_state"], atol=1e-7 ) np.testing.assert_allclose( - padded["present.0.ssm_state"], unpadded["present.0.ssm_state"], atol=1e-7 + padded["present.0.recurrent_state"], + unpadded["present.0.recurrent_state"], + atol=1e-7, ) diff --git a/src/mobius/tasks/_plamo2.py b/src/mobius/tasks/_plamo2.py index 571ff931f..144ae4edd 100644 --- a/src/mobius/tasks/_plamo2.py +++ b/src/mobius/tasks/_plamo2.py @@ -74,7 +74,7 @@ def build(self, module: nn.Module, config: BaseModelConfig) -> ModelPackage: [batch, config.mamba_inner_size, config.mamba_d_conv - 1], ) state_b = builder.input( - f"past_key_values.{layer}.ssm_state", + f"past_key_values.{layer}.recurrent_state", ir.DataType.FLOAT, [ batch, @@ -130,7 +130,7 @@ def build(self, module: nn.Module, config: BaseModelConfig) -> ModelPackage: config.mamba_d_state, ] ) - names = ("conv_state", "ssm_state") + names = ("conv_state", "recurrent_state") state_b.type = ir.TensorType(ir.DataType.FLOAT) state_a.type = ir.TensorType(config.dtype) builder.add_output(state_a, f"present.{layer}.{names[0]}") @@ -139,15 +139,14 @@ def build(self, module: nn.Module, config: BaseModelConfig) -> ModelPackage: model = _make_model(graph) self._register_functions(model, config) model.metadata_props["mobius.cache_abi"] = ( - "per-layer:attention=key,value;mamba=conv_state,ssm_state-f32" + "per-layer:attention=key,value;mamba=conv_state,recurrent_state-f32" ) model.metadata_props["mobius.state_semantics"] = ( "batch-axis reorder;copy-for-rollback;deterministic replay" ) model.metadata_props["mobius.max_verified_context"] = str(config.attention_window_size) model.metadata_props["mobius.runtime_support"] = ( - "deferred: released ORT GenAI cannot represent heterogeneous recurrent/KV " - "state; see onnxruntime/mobius#605" + "ORT GenAI 0.15.2 state ABI; package requires GQA-specialized attention" ) return ModelPackage({"model": model}, config=config) diff --git a/testdata/cases/causal-lm/plamo2-1b.yaml b/testdata/cases/causal-lm/plamo2-1b.yaml index cfd0316d1..0d24ec342 100644 --- a/testdata/cases/causal-lm/plamo2-1b.yaml +++ b/testdata/cases/causal-lm/plamo2-1b.yaml @@ -15,4 +15,4 @@ generation: max_new_tokens: 20 do_sample: false -notes: "The pinned source checkpoint and public F32 GGUF pass raw ORT prefill, all cached decode steps, rollback/replay, and exact 20-token generation. ORT GenAI remains deferred because 0.15.2 cannot bind the heterogeneous recurrent/KV state ABI." +notes: "The pinned source checkpoint and public F32 GGUF pass raw ORT prefill, all cached decode steps, rollback/replay, and exact 20-token generation. ORT GenAI 0.15.2 also passes exact token-ID generation through its sparse KV+recurrent ABI on a CPU-specialized GQA package; unspecialized Attention plus recurrent state is rejected during package config generation. Text tokenization remains deferred because the official custom tokenizer has no identity-preserving tokenizer.json representation." diff --git a/testdata/evidence/causal-lm/plamo2-1b-real.json b/testdata/evidence/causal-lm/plamo2-1b-real.json index 75cf4914b..a8402669d 100644 --- a/testdata/evidence/causal-lm/plamo2-1b-real.json +++ b/testdata/evidence/causal-lm/plamo2-1b-real.json @@ -93,7 +93,7 @@ ], "sha256": "5a9a6638087d4b76fff79e102e665fb963e505cd17812f909cfc33e42a29b640", "state_interfaces": 16, - "state_abi": "even layers: conv_state+ssm_state; odd layers: key+value" + "state_abi": "even layers: conv_state+recurrent_state (semantic SSM state); odd layers: key+value" }, "independent_reference": { "implementation": "pinned HuggingFace remote model with upstream pure-PyTorch causal_conv1d_ref and selective_state_update_ref equivalents", @@ -209,18 +209,46 @@ }, "ort_genai": { "released_version": "0.15.2", + "released_source_tree": "ed5f4e87147731e5b07810f9f5c90103b3603cdf", "generated_model_type": "decoder", - "generated_state_schema": "homogeneous past_key_names/past_value_names templates", + "attention_graph": "CPU-specialized export with all eight attention layers lowered to GroupQueryAttention; unspecialized Attention plus recurrent state is rejected by package config generation", + "generated_state_schema": "sparse global layer indices with key/value on odd layers and conv_state/recurrent_state on even layers", "model_load": "passed with type=decoder and num_hidden_layers=16 discovered from sparse global state indices", - "tokenizer_load": "failed: the exact packaged tokenizer.jsonl format is not accepted by ORT GenAI 0.15.2", - "generation_with_prompt_tokens": "failed: Model input was not found: past_key_values.0.recurrent_state", + "state_contract_proof": "v0.15.2 src/models/kv_cache.cpp independently discovers sparse key inputs; src/models/recurrent_state.cpp independently discovers sparse conv_state inputs and derives paired recurrent_state names", + "generation_with_prompt_tokens": "passed: [1,6721] generated the exact verified 20-token continuation", + "generated_tokens": [ + 44, + 2268, + 1101, + 1080, + 1087, + 8789, + 45114, + 2002, + 3104, + 1131, + 1385, + 1031, + 73, + 39, + 109, + 2569, + 3648, + 353, + 32, + 79135 + ], + "tokenizer_source_revision": "92c75fd6eea9018bcb9c33ee8921589febe071fa", + "tokenizer_jsonl_sha256": "9fa8276a87dd2440bb9496f1d569c05ded6274fd2854de802e12dc6545804b95", + "tokenizer_identity": "all 100000 ordered tokens, token types, float32 scores, and special IDs 0/1/2/3 match the pinned GGUF exactly; the official tokenizer has no normalizer or pre-tokenizer", + "tokenizer_load": "deferred with exact rejection: ORT GenAI 0.15.2 requires tokenizer.json and reports Invalid file; a standard Unigram(byte_fallback=true) conversion is not identical because official encoding of literal <0xFA> is [60,48,120,5109,62] while the standard tokenizer emits byte token [250]", "upstream_main_commit": "a8e0fdf81b061e67c1c3f9485bfdc06735ccd473", "upstream_main_install": "not feasible from the Python source package: pip built UNKNOWN-0.0.0 without an onnxruntime_genai module", - "promotion_decision": "blocked: no released heterogeneous-state schema, no compatible tokenizer artifact, and no successful end-to-end generation" + "promotion_decision": "token-ID generation passed; text-tokenizer entry points remain deferred because no identity-preserving tokenizer.json representation was proven" }, "l4_l5": { "raw_ort_parity": "passed", "generated_tokens_committed": true, - "ort_genai_e2e": "deferred because released ORT GenAI cannot represent the alternating state ABI" + "ort_genai_e2e": "passed on the CPU-specialized GQA package with exact prompt token IDs; text tokenization remains deferred by the proven tokenizer-format mismatch" } } diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 9244cde11..e634fbe74 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -288,8 +288,9 @@ def test_graph_builds_without_weights(self, model_type: str, config_overrides: d assert f"present.{i}.conv_state" in output_names, ( f"Missing present.{i}.conv_state" ) - assert f"present.{i}.ssm_state" in output_names, ( - f"Missing present.{i}.ssm_state" + state_name = "recurrent_state" if model_type == "plamo2" else "ssm_state" + assert f"present.{i}.{state_name}" in output_names, ( + f"Missing present.{i}.{state_name}" ) elif ltype == "conv": assert f"present.{i}.conv_state" in output_names, (