diff --git a/.agents/skills/mobius-onnx-export-gotchas/SKILL.md b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md new file mode 100644 index 00000000..7b44c17d --- /dev/null +++ b/.agents/skills/mobius-onnx-export-gotchas/SKILL.md @@ -0,0 +1,200 @@ +--- +name: mobius-onnx-export-gotchas +description: Use when building/exporting ONNX models with the `mobius build` CLI (especially Phi-3 / Phi-3.5 or any model with `--execution-provider cuda` GQA fusion and/or `--static-cache`). Covers the current CLI syntax, the dtype flag values, the GQA-vs-static-cache interaction, how to verify fp16 GQA exports load in onnxruntime (the historical packed-QKV FLOAT32 load bug is fixed by the fp16 GQA fold-fix), and why fp16 GQA exports need VALUE-based weight checks (corr≈1.0 / norm), not just initializer count/dtype, to catch silently-zeroed packed-QKV weights. +--- + +# mobius ONNX export gotchas + +## 1. CLI syntax (editable repo differs from older docs) +`mobius build` requires `--model ` and takes the **output dir as a POSITIONAL** arg. +There is **no `-o` flag for `build`** (`-o` exists only on `build-gguf`). + +```bash +mobius build --model microsoft/Phi-3.5-mini-instruct \ + --dtype f16 --execution-provider cuda \ + --external-data onnx --trust-remote-code \ + /path/to/output_dir +``` + +- `--dtype` choices: `f16`/`float16`, `bf16`/`bfloat16`, `f32`/`float32`. **`fp16` is INVALID.** +- `--execution-provider` is an alias of `--ep`. `cuda` + fp16/bf16 triggers GQA fusion; + `default` keeps plain ONNX `Attention`. + +## 2. `--static-cache` is incompatible with GQA fusion +`--static-cache` wraps each attention with `TensorScatter` (in-place KV cache for the **ONNX Attention** +op). That breaks the pattern the GQA rewrite matches, so combining +`--execution-provider cuda --static-cache` yields **0 GroupQueryAttention + N Attention + 2N TensorScatter** +(mobius prints: "GQA fusion expected … but found 0 GroupQueryAttention and N Attention nodes"). + +- **GQA model:** `--execution-provider cuda` **alone**. GQA's shared KV buffer + (`past_present_share_buffer`) is enabled at **runtime** via IO-binding past & present to the same + OrtValue — NOT via `--static-cache`. +- **ONNX-Attention + in-place cache:** `--execution-provider default --static-cache --max-seq-len N`. + +## 3. FIXED: fp16 GQA export previously left packed-QKV weights as FLOAT32 → model wouldn't load +**Status: fixed (the fp16 GQA fold-fix).** Native fp16 Phi-3.5 GQA export now loads directly in the ORT +CUDA EP with **no manual post-cast** (32 GroupQueryAttention nodes, all-fp16 initializers). If you are on +that commit or later, you should not hit this — skip to the verification snippet below. The history is +kept here because old artifacts exported before the fix still carry fp32 packed weights. + +### Symptom (pre-fix) +For an fp16 GQA export, a folded per-layer packed QKV weight +(`..q_proj.weight__k_proj.weight__v_proj.weight__axis_0__concat`) was emitted as **FLOAT32**, while its +MatMul's other input was fp16. onnxruntime then rejected the model at load on both CPU and CUDA EPs: + +``` +Type Error: Type parameter (T) of Optype (MatMul) bound to different types +(tensor(float16) and tensor(float)) in node (node_MatMul_*) +``` + +You'd also see at save time: `The value type for shape [H, 3H] is not known. Skipping serialization`. + +### Root cause +`_cast_module_dtype` casts module params to fp16, but the resulting initializer `Value`s lose their +declared `.dtype` (it becomes `None`) while their `const_value` stays fp16. The fold passes +`FoldConcatInitializersPass` (`src/mobius/_passes/_fold_concat.py`) and `FoldTransposedInitializerPass` +(`src/mobius/_passes/_fold_transpose.py`) then defaulted the folded initializer's dtype to `FLOAT`, +serializing the packed QKV / transposed weights as fp32. + +### The fix +A shared helper `initializer_dtype()` (`src/mobius/_passes/_dtype_utils.py`) resolves the effective dtype +from the declared type, **falling back to `const_value` when the type annotation was dropped** (preferring +the data dtype and warning on stale-metadata disagreement). Both fold passes use it to stamp the correct +dtype on the new initializer's `TensorType` and `LazyTensor`, and `FoldConcatInitializersPass` now also +skips folding before weights are loaded (mirroring `FoldTransposedInitializerPass`). A regression test +loads the fp16 GQA export in the ORT CPU EP to lock this in. + +### Verify (still worth running on any fp16 build) +```python +import onnx +m = onnx.load("model.onnx", load_external_data=False) +fp32 = [i.name for i in m.graph.initializer if i.data_type == onnx.TensorProto.FLOAT] +print(len(fp32), "FLOAT32 initializers (should be 0 for fp16)") +``` + +### Convention (prevents the whole class from reappearing) +Any pass that **materializes a new initializer** must resolve its dtype via +`initializer_dtype()` (`src/mobius/_passes/_dtype_utils.py`), **never** `value.dtype or ir.DataType.FLOAT`. +The bug class originates in `_cast_module_dtype` dropping a `Value`'s declared `.dtype` (→ `None`) while its +`const_value` stays fp16; a bare `.dtype or FLOAT` fallback then silently mis-types the result as fp32. Fold +passes (`_fold_concat.py`, `_fold_transpose.py`) already follow this; mirror it in any future +initializer-producing pass. Siblings still reading `.dtype` directly remain exposed — a follow-up should +grep `_passes/` for `.dtype or ir.DataType` and consider re-stamping the type in `_cast_module_dtype` to kill +the class at source. + +### Salvaging a stale pre-fix artifact (only if re-exporting is not an option) +Prefer re-exporting on the fixed code. If you must repair an old model, cast its FLOAT32 initializers to +fp16 and re-save. **Gotcha when re-saving with external data:** if you save with `location="X.data"` and +then rename the file, the references inside `model.onnx` still point to `X.data`. Either save directly +with `location="model.onnx.data"`, or rewrite each initializer's `external_data` `location` entry. + +```python +import onnx, numpy as np +from onnx import numpy_helper, TensorProto +m = onnx.load("model.onnx", load_external_data=True) +for init in m.graph.initializer: + if init.data_type == TensorProto.FLOAT: + arr = numpy_helper.to_array(init).astype(np.float16) + init.CopyFrom(numpy_helper.from_array(arr, init.name)) +onnx.save(m, "model.onnx", save_as_external_data=True, all_tensors_to_one_file=True, + location="model.onnx.data", size_threshold=1024, convert_attribute=False) +``` + +## 4. Always validate the export in ORT before profiling +Load the model on `CUDAExecutionProvider` and run one prefill + one decode `session.run`. Confirm: +(a) the expected attention op (`com.microsoft::GroupQueryAttention` vs `ai.onnx::Attention`), +(b) finite fp16 logits, (c) no FLOAT32 initializers for an fp16 build. + +These checks are **necessary but NOT sufficient** for a fp16 GQA export — see §5. A model can pass all +three and still have silently-zeroed packed-QKV weights. + +## 5. Verifying a fp16 GQA export: use VALUE-based weight checks, NOT initializer count/dtype +**A fp16 GQA export can be all-fp16, right-count, and still all-zeros — only a corr≈1.0 / norm≈126 VALUE +check on the packed QKV proves the weights are real.** + +### Symptom +The GQA model loads cleanly (32 `GroupQueryAttention` nodes, all-fp16, finite logits) but generates +garbage (e.g. `holdou_(...artersarters`). Prefill logits come out ~3× the reference scale, with +`max|Δlogit|` ~50+ versus the reference. + +### Root cause +The packed-QKV initializer is `Transpose(Concat(q, k, v, axis=0))`. If the fold passes +(`FoldConcatInitializersPass` / `FoldTransposedInitializerPass`) leave the packed-Concat output dtype +UNKNOWN / defaulted-to-fp32 while the data is fp16, the serializer **skips** it and it loads as +**near-zero** — the weights are silently dead. (This is the §3 failure mode; the upstream fix in +The fp16 GQA fold-fix stamps the fp16 dtype at the fold-pass source. A post-hoc cast is NOT a fix — it re-corrupts.) + +### Why count/dtype checks fail (the trap) +The BROKEN export and the FIXED export can have the **same initializer count and the same fp16/fp32 dtype +ratio**, so neither is a validity signal. Worse, the fp16-init count is **not even stable across fixes** +— on Phi-3.5 it moved from ~293 down to ~197 (an unstripped intermediate carries the packed-QKV plus the +now-dead unpacked q/k/v source initializers; a safe dead-weight strip then removes the ~96 dead pre-pack +inits), with no bearing on correctness. Note the OLD broken export was *also* 197 fp16, so even a "right" +final count proves nothing. Counting initializers or checking "0 fp32 / all fp16" does **not** distinguish a +healthy model from a zeroed-weight one. §4(c) alone will pass a dead model. **Never gate on the count; +use the VALUE gate below.** + +### Canonical verification (load-bearing, not optional) +VALUE-based per-slice check on each packed-QKV initializer against its source q/k/v weights: +- per-slice correlation **≈ 1.000** (broken ≈ 0.000), AND +- packed-QKV L2 norm **≈ 126.6** at layer 0 / mean(|abs|) **≈ 0.015** (broken ≈ 0.80 / ≈ 5e-6). + +> ⚠️ **Use mean-of-ABS or norm — NEVER the signed mean.** The good model's *signed* mean is ~2.6e-6 +> (near zero, because the weights are symmetric ±), which coincidentally looks just like the broken +> model's mean(|abs|) ~5e-6. Checking signed mean would **falsely flag the good model as broken** — this +> exact confusion has already caused a false alarm in this crew. Valid discriminators: mean(|abs|) +> (good ≈ 0.015 vs broken ≈ 5e-6) or L2 norm (good ≈ 126.6 vs broken ≈ 0.80). + +Plus an end-to-end next-token greedy-argmax parity check vs the `attn_dynamic` reference (expect +**~19–20 / 20**). Isolated single-token divergences are fp16 dead-ties (reference top1−top2 gap = 0.0000), +not bugs. Optional hardening: assert **0 unused initializers** and that all N packed-QKV initializers are +present, to catch dead-weight OVER-stripping. + +QA's `gqa_weight_integrity_gate.py` (`--self-check --strip-audit --scan-all`, per-layer corr/norm) +implements exactly this gate. + +## 6. FIXED: GQA `present.*` KV-cache outputs declared the wrong `head_dim` +**Status: fixed (the GQA present-KV shape fix).** A native fp16 GQA export now declares +`present.{i}.{key,value}` with the correct `head_dim`, symmetric to its `past_key_values.{i}.*` inputs. + +### Symptom (pre-fix) +The graph **output** `present.{i}.key/value` declared the wrong `head_dim` (e.g. `32` instead of the real +`96` on Phi-3.5) while the matching `past_key_values.{i}.*` **input** was correct (`96`). At load ORT logged +(once per key+value per layer — 64 on Phi-3.5): + +``` +[W ...MergeShapeInfo] Error merging shape info for output. 'present.0.key' +source:{-1,32,-1,96} target:{-1,32,-1,32}. Falling back to lenient merge. +``` + +Runtime still produced correct (96-wide) arrays via lenient merge, but any consumer that **trusts declared +shapes** (e.g. `onnxruntime-genai`) would see inconsistent past-vs-present KV cache types. + +### Root cause +`GroupQueryAttention`'s contrib-op shape inference mis-derives the present `head_dim` (it does **not** +reproduce on the plain `Attention` op, which infers correctly). `_register_kv_cache_outputs` +(`src/mobius/tasks/_cache_utils.py`) added the present outputs with **no explicit shape**, so the buggy +inference won. + +### The fix +`_register_kv_cache_outputs` now opt-in **stamps** `present.{i}.{key,value}` shape+dtype symmetric to the +past inputs when the caller passes `batch`/`num_kv_heads`/`key_head_dim`/`value_head_dim`/`total_seq_len`/ +`dtype` (wired from `_causal_lm.py`). Omitting them preserves inference-only behavior, so the other ~10 +callers are unaffected. The stamp survives `SymbolicShapeInferencePass` (policy `refine` only tightens +unknown dims; it won't replace a concrete `96` with a conflicting `32`). + +### Verify +```python +import onnx +m = onnx.load("model.onnx", load_external_data=False) +d = lambda vi: [(x.dim_param or x.dim_value) for x in vi.type.tensor_type.shape.dim] +o = {v.name: v for v in m.graph.output} +print("present.0.key:", d(o["present.0.key"])) # head_dim must equal the past input's (e.g. 96, NOT 32) +``` + +### Known remaining (separate, pre-existing, harmless) +ORT still logs ~32 `Error merging shape info ... source:{-1,-1,3072} target:{-1,-1,1024}` warnings on the +GQA op's **internal hidden-state output** value_info (`v_*.GroupQueryAttention_*_0`, `1024`=32×32 vs the +correct `3072`=32×96). That value is **not** a declared graph I/O — runtime is correct and `onnxruntime-genai` +does not trust it — so it does not bite shape-trusting consumers the way the present-output bug did. Tracked +as a follow-up in the GQA rewrite emission path (not the KV-cache output path). diff --git a/CHANGELOG.md b/CHANGELOG.md index 83edfb24..318da777 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,57 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### KV-cache present-shape: fail-closed on partial parameter sets + +#### Fixed + +- `_register_kv_cache_outputs` now **raises `ValueError`** when given a partial + set of present-shape parameters (1–5 of the six `batch`, `num_kv_heads`, + `key_head_dim`, `value_head_dim`, `total_seq_len`, `dtype`) instead of logging + a warning and proceeding. A partial set is always a wiring slip with no + legitimate use; the previous fail-open shipped a structurally-wrong model + (mis-derived `GroupQueryAttention` present `head_dim`) with only a log line. + Passing all six (stamp) or none (infer) is unaffected. (closes #341) + +--- + +### fp16 GQA Export Fix + +#### Fixed + +- Native fp16 GroupQueryAttention exports (e.g. `microsoft/Phi-3.5-mini-instruct` + with `--dtype f16 --execution-provider cuda`) no longer emit fp32 packed-QKV / + transposed weights. Previously the fold passes (`FoldConcatInitializersPass`, + `FoldTransposedInitializerPass`) defaulted a folded initializer's dtype to + `FLOAT` when the source `Value`'s declared type had been dropped during fp16 + casting, producing a model onnxruntime rejected at load with a + `MatMul` type-parameter error (`tensor(float16)` vs `tensor(float)`) on both + CPU and CUDA EPs. A new `mobius._passes._dtype_utils.initializer_dtype()` + helper now resolves the effective dtype from `const_value` when the type + annotation is missing, so fp16 GQA models load directly with no manual + post-cast. + +--- + +### GQA Present KV-Cache Shape Fix + +#### Fixed + +- GroupQueryAttention exports now declare correct `present.{i}.key` / + `present.{i}.value` graph-output shapes and dtype. The GQA contrib op's shape + inference mis-derived the present KV `head_dim` (e.g. 32 instead of 96 on + `microsoft/Phi-3.5-mini-instruct`), so the present KV-cache outputs declared a + `head_dim` inconsistent with the (correct) `past_key_values` inputs. ORT logged + `Error merging shape info ... lenient merge` (64 warnings on Phi-3.5) and any + consumer that chains `present` → `past` and trusts declared shapes (e.g. + `onnxruntime-genai`) saw mismatched past-vs-present KV cache types. This is a + metadata / declared-shape correction only — runtime numerics are unchanged + (weights byte-identical, next-token parity 20/20). `_register_kv_cache_outputs` + now stamps the present KV outputs symmetric to the past inputs. Affects + GQA-fusion packed-QKV exports (Phi-3.5, Llama-3.2, Qwen2, Mistral, Phi-3-GQA). + +--- + ### WebGPU Shape Op Support #### Changed diff --git a/src/mobius/_passes/_dtype_utils.py b/src/mobius/_passes/_dtype_utils.py new file mode 100644 index 00000000..a05a5c51 --- /dev/null +++ b/src/mobius/_passes/_dtype_utils.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Shared dtype helpers for graph passes that materialize new initializers. + +Passes such as :class:`~mobius._passes.FoldConcatInitializersPass` and +:class:`~mobius._passes.FoldTransposedInitializerPass` pre-compute new +initializers from existing ones. They must stamp the *correct* dtype on the +result, otherwise an fp16 model can silently end up with fp32 weights that +onnxruntime rejects at load time (a MatMul binding fp16 and fp32 to the same +type parameter ``T``). +""" + +from __future__ import annotations + +import onnx_ir as ir + + +def initializer_dtype(value: ir.Value) -> ir.DataType | None: + """Return the effective dtype of an initializer ``value``. + + Uses the value's declared ``type`` dtype, but falls back to the dtype of its + ``const_value`` when the type annotation is missing. + + Graph building can drop the declared ``type`` on an initializer while its + actual tensor data (``const_value``) still carries the correct dtype. In + that situation, defaulting to ``ir.DataType.FLOAT`` would emit fp32 weights + into an otherwise fp16 model. Reading the dtype from ``const_value`` keeps + folded initializers consistent with the weights they are derived from. + + When both a declared type and ``const_value`` are present but **disagree**, + this raises :class:`ValueError`: a declared type that contradicts the + serialized data is corrupt metadata with no legitimate use, so fail closed + (consistent with the fail-closed contract elsewhere in the export pipeline) + rather than silently picking one and shipping a structurally-wrong model. + The declared-is-``None`` fallback — the path this helper exists for — is + unaffected and never raises. + + Returns ``None`` only when neither the declared type nor ``const_value`` is + available; callers decide on a final fallback. + """ + declared = value.dtype + const_dtype = value.const_value.dtype if value.const_value is not None else None + + if declared is not None and const_dtype is not None and declared != const_dtype: + raise ValueError( + f"Initializer {value.name!r} declares dtype {declared} but its " + f"const_value data is {const_dtype}. A declared type that " + f"contradicts the serialized data indicates corrupt initializer " + f"metadata." + ) + if declared is not None: + return declared + return const_dtype diff --git a/src/mobius/_passes/_dtype_utils_test.py b/src/mobius/_passes/_dtype_utils_test.py new file mode 100644 index 00000000..e3be7805 --- /dev/null +++ b/src/mobius/_passes/_dtype_utils_test.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the shared initializer dtype helper.""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import pytest + +from mobius._passes._dtype_utils import initializer_dtype + + +def _value(dtype: ir.DataType | None, const: np.ndarray | None) -> ir.Value: + # Built via the ir.Value constructor rather than the ir.val() factory: these + # fixtures deliberately construct degenerate initializers (a dropped declared + # type alongside a const_value, and a declared type that contradicts the + # const_value dtype) that ir.val() validates against and refuses to build. + return ir.Value( + name="w", + type=ir.TensorType(dtype) if dtype is not None else None, + const_value=ir.tensor(const) if const is not None else None, + ) + + +class TestInitializerDtype: + def test_uses_declared_dtype_when_present(self): + v = _value(ir.DataType.FLOAT16, np.ones((2,), np.float16)) + assert initializer_dtype(v) == ir.DataType.FLOAT16 + + def test_falls_back_to_const_value_when_declared_missing(self): + """The core fix: a dropped declared type must not hide the real dtype.""" + v = _value(None, np.ones((2,), np.float16)) + assert v.dtype is None + assert initializer_dtype(v) == ir.DataType.FLOAT16 + + def test_raises_on_dtype_contradiction(self): + """Reject values whose declared dtype contradicts the serialized data. + + Such a value is corrupt metadata and must fail closed rather than + silently pick one dtype. + """ + v = _value(ir.DataType.FLOAT, np.ones((2,), np.float16)) + with pytest.raises(ValueError) as excinfo: + initializer_dtype(v) + message = str(excinfo.value) + assert "FLOAT" in message and "FLOAT16" in message + assert "w" in message + + def test_returns_none_when_nothing_available(self): + v = _value(None, None) + assert initializer_dtype(v) is None diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py index b8f2eed9..7035f1e9 100644 --- a/src/mobius/_passes/_fold_concat.py +++ b/src/mobius/_passes/_fold_concat.py @@ -27,6 +27,8 @@ import numpy as np import onnx_ir as ir +from mobius._passes._dtype_utils import initializer_dtype + logger = logging.getLogger(__name__) @@ -63,6 +65,19 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: if not all(v is not None and v.is_initializer() for v in inputs): continue + # All inputs must have loaded weights. Folding before weights are + # loaded would force the dtype to default to FLOAT and bake a wrong + # type into the packed initializer. Mirror FoldTransposedInitializer's + # guard and skip with a warning so the Concat survives for a later run. + if any(v.const_value is None for v in inputs): # type: ignore[union-attr] + logger.warning( + "FoldConcatInitializersPass: skipping Concat %r — an input " + "initializer has no const_value (pass ran before weights " + "were loaded, or a weight is missing).", + node.name, + ) + continue + axis_attr = node.attributes.get("axis") axis: int = axis_attr.value if axis_attr is not None else 0 @@ -79,8 +94,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: # Require uniform dtype — mixed-dtype concat is unusual and likely # a modelling error; skip and warn rather than silently produce - # a wrong result. - dtypes = [v.dtype for v in inputs] # type: ignore[union-attr] + # a wrong result. Resolve each dtype from the declared type or, when + # that is missing, from the loaded ``const_value`` so an fp16 weight + # whose type annotation was dropped is not mistaken for fp32. + dtypes = [initializer_dtype(v) for v in inputs] # type: ignore[union-attr] if len(set(dtypes)) > 1: logger.warning( "FoldConcatInitializersPass: skipping Concat with mixed dtypes %s" @@ -89,6 +106,10 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: ) continue + # Authoritative dtype for the packed initializer. Falls back to FLOAT + # only when neither the declared type nor const_value is available. + packed_dtype = dtypes[0] or ir.DataType.FLOAT + # Build a name for the packed initializer from the input names and axis. # The name encodes both the ordered input names and the axis, so two Concat # nodes with the same inputs in the same order along the same axis will @@ -104,7 +125,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: if packed_name in model.graph.initializers: existing_val = model.graph.initializers[packed_name] out_val.replace_all_uses_with(existing_val, replace_graph_outputs=True) - model.graph.remove(node) + # safe=True detaches the node from its inputs before removal so + # the folded q/k/v initializers' use lists are cleared. Otherwise + # `inp.uses()` still points at the removed Concat and the dead + # pre-pack weights survive RemoveUnusedNodesPass (~1.8 GB of + # orphaned weights serialized into the fp16 GQA model). + model.graph.remove(node, safe=True) folded_count += 1 modified = True @@ -117,7 +143,12 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: ) continue - new_val = ir.Value(name=packed_name, shape=out_shape, type=out_val.type) + # Stamp the resolved dtype on the new initializer's type so the + # declared type, the LazyTensor, and the materialized data all agree + # (out_val.type may be missing after stage-2 rewrites). + new_val = ir.Value( + name=packed_name, shape=out_shape, type=ir.TensorType(packed_dtype) + ) captured_inputs = list(inputs) # capture for closure @@ -133,13 +164,17 @@ def _make_packed( arrays.append(v.const_value.numpy()) return ir.tensor(np.concatenate(arrays, axis=ax)) - dtype = inputs[0].dtype or ir.DataType.FLOAT # type: ignore[union-attr] - new_val.const_value = ir.LazyTensor(_make_packed, dtype=dtype, shape=out_shape) + new_val.const_value = ir.LazyTensor( + _make_packed, dtype=packed_dtype, shape=out_shape + ) model.graph.initializers[new_val.name] = new_val out_val.replace_all_uses_with(new_val, replace_graph_outputs=True) - model.graph.remove(node) + # safe=True detaches the node from its inputs before removal so the + # folded q/k/v initializers' use lists are cleared; otherwise the + # dead pre-pack weights survive RemoveUnusedNodesPass. + model.graph.remove(node, safe=True) folded_count += 1 modified = True diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py index 492538d6..d8bb0416 100644 --- a/src/mobius/_passes/_fold_concat_test.py +++ b/src/mobius/_passes/_fold_concat_test.py @@ -19,10 +19,11 @@ def _make_concat_model( """Build a model with a Concat node over initializers.""" init_vals: list[ir.Value] = [] for i, arr in enumerate(arrays): + tensor = ir.tensor(arr) v = ir.Value(name=f"init_{i}") v.shape = ir.Shape(list(arr.shape)) - v.dtype = ir.DataType.FLOAT - v.const_value = ir.tensor(arr) + v.dtype = tensor.dtype + v.const_value = tensor init_vals.append(v) inputs_to_concat: list[ir.Value] = list(init_vals) @@ -44,7 +45,7 @@ def _make_concat_model( cat_shape[axis] = sum(a.shape[axis] for a in arrays) if not add_dynamic_input: out.shape = ir.Shape(cat_shape) - out.dtype = ir.DataType.FLOAT + out.dtype = ir.tensor(arrays[0]).dtype graph_inputs = [ir.Value(name="x")] if add_dynamic_input: @@ -124,6 +125,61 @@ def test_dynamic_input_not_folded(self): result = FoldConcatInitializersPass()(model) assert not result.modified + def test_folded_inputs_detached_so_dce_strips_dead_weights(self): + """After folding, the source initializers must be detached and DCE-able. + + FoldConcat removes the Concat node; if it does not detach the node from + its inputs (``safe=True``), the source q/k/v initializers keep a stale + use pointing at the removed node, so RemoveUnusedNodesPass cannot strip + them. For fp16 GQA that leaves ~1.8 GB of dead pre-pack weights in the + exported model. + """ + from onnx_ir.passes import common as common_passes + + a = np.ones((4, 8), dtype=np.float16) + b = np.ones((4, 8), dtype=np.float16) * 2 + model, init_vals = _make_concat_model([a, b], axis=0) + # The Concat output is the only graph output; route it through a MatMul + # so the packed initializer stays live while the sources become dead. + packed_out = model.graph.outputs[0] + hidden = ir.Value( + name="hidden", + shape=ir.Shape(["N", 8]), + type=ir.TensorType(ir.DataType.FLOAT16), + ) + matmul = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1) + mm_out = matmul.outputs[0] + mm_out.shape = ir.Shape(["N", 8]) + mm_out.dtype = ir.DataType.FLOAT16 + model.graph.append(matmul) + model.graph.inputs.append(hidden) + model.graph.outputs[0] = mm_out + + FoldConcatInitializersPass()(model) + + # Sources must be detached (zero uses) so DCE can remove them. + for v in init_vals: + assert len(list(v.uses())) == 0, ( + f"{v.name} still has uses after fold — node not detached" + ) + + common_passes.RemoveUnusedNodesPass()(model) + remaining = set(model.graph.initializers) + assert "init_0" not in remaining and "init_1" not in remaining, ( + f"Dead pre-pack weights survived DCE: {remaining}" + ) + assert "init_0__init_1__axis_0__concat" in remaining, ( + f"the live packed result must NOT be stripped by DCE: {remaining}" + ) + # The survived packed weight must retain its exact values, not just its + # name — guards against a future DCE that mutates retained tensors. + packed = model.graph.initializers["init_0__init_1__axis_0__concat"] + np.testing.assert_array_equal( + packed.const_value.numpy(), + np.concatenate([a, b], axis=0), + err_msg="packed-QKV values corrupted by DCE", + ) + def test_uses_lazy_tensor(self): """The packed initializer wraps sources in a LazyTensor. @@ -467,6 +523,56 @@ def test_concat_with_no_inputs_skipped(self): result = FoldConcatInitializersPass()(model) assert not result.modified + def test_packed_dtype_follows_const_value_when_declared_dtype_missing(self): + """Regression: fp16 weights with an unset declared dtype must fold to fp16. + + After ``_cast_module_dtype`` casts parameters to fp16, the resulting + initializer Values keep an fp16 ``const_value`` but lose their declared + ``.dtype`` (it becomes ``None``). The pass must resolve the packed dtype + from ``const_value`` rather than defaulting to FLOAT, otherwise GQA's + packed-QKV weight is serialized as fp32 and ORT rejects the model with a + fp16/fp32 MatMul type-mismatch error. + """ + q = ir.Value(name="q_weight") + q.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16)) + k = ir.Value(name="k_weight") + k.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16)) + v = ir.Value(name="v_weight") + v.const_value = ir.tensor(np.ones((4, 3), dtype=np.float16)) + # Critically: do NOT set .dtype — declared dtype stays None. + assert q.dtype is None and k.dtype is None and v.dtype is None + + concat_node = ir.Node( + "", + "Concat", + inputs=[q, k, v], + attributes=[ir.Attr("axis", ir.AttributeType.INT, 0)], + num_outputs=1, + ) + graph = ir.Graph( + inputs=[ir.Value(name="x")], + outputs=[concat_node.outputs[0]], + nodes=[concat_node], + name="test_graph", + opset_imports={"": 20}, + ) + graph.register_initializer(q) + graph.register_initializer(k) + graph.register_initializer(v) + model = ir.Model(graph, ir_version=10) + + result = FoldConcatInitializersPass()(model) + assert result.modified + + packed = model.graph.initializers["q_weight__k_weight__v_weight__axis_0__concat"] + assert packed.dtype == ir.DataType.FLOAT16, ( + f"Packed initializer declared dtype should be FLOAT16, got {packed.dtype}" + ) + assert packed.const_value.dtype == ir.DataType.FLOAT16, ( + f"Packed LazyTensor dtype should be FLOAT16, got {packed.const_value.dtype}" + ) + assert packed.const_value.numpy().dtype == np.float16 + def test_shape_computed_from_inputs_when_output_shape_missing(self): """When the Concat output has no shape, shape is inferred from input shapes.""" a = np.ones((4, 3), dtype=np.float32) @@ -512,3 +618,275 @@ def test_shape_computed_from_inputs_when_output_shape_missing(self): # Shape should be computed from inputs: (4,3) + (4,3) along axis 0 → (8,3). packed = model.graph.initializers[packed_name] assert list(packed.shape) == [8, 3] + + +class TestFoldConcatOrtLoad: + """End-to-end regression: a folded fp16 GQA-style graph must load in ORT. + + Reproduces the original failure where the packed-QKV initializer was + serialized as fp32 while the rest of the graph was fp16, causing ORT to + reject the model with:: + + Type parameter (T) of Optype (MatMul) bound to different types + (tensor(float16) and tensor(float)) + + The failure occurs on the CPU EP too, so no GPU is needed. + """ + + def test_folded_fp16_concat_matmul_loads_in_ort(self, tmp_path): + import onnxruntime as ort + + # Three fp16 weights with UNSET declared dtype (as after _cast_module_dtype), + # but with shape retained so they serialize before they are pruned. + q = ir.Value(name="q_weight", shape=ir.Shape([8, 4])) + q.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16)) + k = ir.Value(name="k_weight", shape=ir.Shape([8, 4])) + k.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16)) + v = ir.Value(name="v_weight", shape=ir.Shape([8, 4])) + v.const_value = ir.tensor(np.random.randn(8, 4).astype(np.float16)) + + concat_node = ir.Node( + "", + "Concat", + inputs=[q, k, v], + attributes=[ir.Attr("axis", ir.AttributeType.INT, 1)], + num_outputs=1, + ) + packed_out = concat_node.outputs[0] # shape (8, 12), packed QKV weight + packed_out.shape = ir.Shape([8, 12]) + + # hidden @ packed_qkv : (N, 8) @ (8, 12) -> (N, 12), all fp16. + hidden = ir.Value( + name="hidden", shape=ir.Shape(["N", 8]), type=ir.TensorType(ir.DataType.FLOAT16) + ) + matmul_node = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1) + out = matmul_node.outputs[0] + out.shape = ir.Shape(["N", 12]) + out.dtype = ir.DataType.FLOAT16 + + graph = ir.Graph( + inputs=[hidden], + outputs=[out], + nodes=[concat_node, matmul_node], + name="qkv_matmul", + opset_imports={"": 20}, + ) + graph.register_initializer(q) + graph.register_initializer(k) + graph.register_initializer(v) + model = ir.Model(graph, ir_version=10) + + FoldConcatInitializersPass()(model) + + # Drop the now-unused source initializers (a DCE step runs after folding + # in the real export pipeline). Without this they linger as dead weights. + for dead in ("q_weight", "k_weight", "v_weight"): + del model.graph.initializers[dead] + + model_path = tmp_path / "fp16_qkv.onnx" + ir.save(model, model_path) + + # Before the fix this raises a fp16/fp32 MatMul type-mismatch at load. + sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + feed = {"hidden": np.random.randn(2, 8).astype(np.float16)} + (result,) = sess.run(None, feed) + assert result.shape == (2, 12) + assert result.dtype == np.float16 + + def test_folded_fp16_packed_qkv_values_survive_serialization(self, tmp_path): + """Value gate: packed-QKV must keep its fp16 *values* through serialize→reload. + + The other fold tests assert the packed value in memory and that ORT can + load+run the model, but none compare the *serialized* packed-QKV weight + to its source q/k/v projections. The original garbage-export bug + corrupted bytes at serialization (fp16 data written under a defaulted + FLOAT32 dtype → near-zero/garbage) — a failure an in-memory + ``const_value`` check cannot see, and which a load+run check misses + because the model still loads and produces a right-shaped fp16 output. + + This round-trips through the production save path (``ir.save`` with + external data, exactly like the real fp16 export's model.onnx + + model.onnx.data) and asserts the reloaded packed weight matches its + sources per-slice, catching a numerically-corrupt pack that still has + the right count and dtype (the corr≈0 / norm≈0.8 failure mode). + """ + import onnxruntime as ort + + rng = np.random.default_rng(0) + # fp16 q/k/v projection weights with UNSET declared dtype (the state left + # by _cast_module_dtype). Sized > 256 bytes so ir.save externalizes them, + # exercising the same external-data path as the real export. Column-concat + # (axis=1) mirrors the real packed-QKV layout [Q | K | V]. + q_arr = rng.standard_normal((16, 16)).astype(np.float16) + k_arr = rng.standard_normal((16, 16)).astype(np.float16) + v_arr = rng.standard_normal((16, 16)).astype(np.float16) + q = ir.Value(name="q_weight", shape=ir.Shape([16, 16])) + q.const_value = ir.tensor(q_arr) + k = ir.Value(name="k_weight", shape=ir.Shape([16, 16])) + k.const_value = ir.tensor(k_arr) + v = ir.Value(name="v_weight", shape=ir.Shape([16, 16])) + v.const_value = ir.tensor(v_arr) + + concat_node = ir.Node( + "", + "Concat", + inputs=[q, k, v], + attributes=[ir.Attr("axis", ir.AttributeType.INT, 1)], + num_outputs=1, + ) + packed_out = concat_node.outputs[0] # (16, 48) packed QKV weight + packed_out.shape = ir.Shape([16, 48]) + + hidden = ir.Value( + name="hidden", shape=ir.Shape(["N", 16]), type=ir.TensorType(ir.DataType.FLOAT16) + ) + matmul_node = ir.Node("", "MatMul", inputs=[hidden, packed_out], num_outputs=1) + out = matmul_node.outputs[0] + out.shape = ir.Shape(["N", 48]) + out.dtype = ir.DataType.FLOAT16 + + graph = ir.Graph( + inputs=[hidden], + outputs=[out], + nodes=[concat_node, matmul_node], + name="qkv_matmul", + opset_imports={"": 20}, + ) + graph.register_initializer(q) + graph.register_initializer(k) + graph.register_initializer(v) + model = ir.Model(graph, ir_version=10) + + FoldConcatInitializersPass()(model) + # DCE step the real export runs after folding. + for dead in ("q_weight", "k_weight", "v_weight"): + del model.graph.initializers[dead] + + # Serialize through the production path: external data, exactly as the + # real fp16 export writes model.onnx + model.onnx.data. This is where the + # original dtype bug corrupted the packed bytes. + model_path = tmp_path / "model.onnx" + ir.save(model, model_path, external_data="model.onnx.data") + assert (tmp_path / "model.onnx.data").exists(), ( + "weights should be externalized, exercising the real export's save path" + ) + + reloaded = ir.load(model_path) + packed_name = "q_weight__k_weight__v_weight__axis_1__concat" + packed = reloaded.graph.initializers[packed_name].const_value.numpy() + assert packed.dtype == np.float16, ( + f"reloaded packed-QKV must serialize as fp16, got {packed.dtype}" + ) + expected = np.concatenate([q_arr, k_arr, v_arr], axis=1) + assert packed.shape == expected.shape + + # Per-slice discriminator (mirrors QA's weight-integrity gate): each q/k/v + # column block of the reloaded pack must match its source projection. A + # serialize-time dtype corruption surfaces here as corr≈0 / norm collapse + # even when the pack's shape and dtype look correct. + blocks = { + "q": (packed[:, 0:16], q_arr), + "k": (packed[:, 16:32], k_arr), + "v": (packed[:, 32:48], v_arr), + } + for name, (got, src) in blocks.items(): + got32 = got.astype(np.float32).ravel() + src32 = src.astype(np.float32).ravel() + corr = np.corrcoef(got32, src32)[0, 1] + assert corr >= 0.99, ( + f"{name}-slice packed-QKV corr={corr:.4f} (<0.99) — corrupted pack" + ) + src_norm = np.linalg.norm(src32) + rel_err = abs(np.linalg.norm(got32) - src_norm) / src_norm + assert rel_err <= 0.02, ( + f"{name}-slice packed-QKV norm rel_err={rel_err:.4f} (>2%) — corrupted pack" + ) + # Non-degenerate magnitude. Use mean of ABSOLUTE values, not signed + # mean: symmetric fp16 weights have a signed mean ~1e-6 that is + # indistinguishable from a broken near-zero tensor, whereas mean|abs| + # separates cleanly (healthy ~0.0x vs unserialized ~1e-6). This also + # catches the near-zero failure mode robustly where corr is undefined + # (a constant/zero slice has zero variance → corrcoef is nan). + mean_abs = float(np.abs(got32).mean()) + assert mean_abs >= 1e-3, ( + f"{name}-slice packed-QKV mean|abs|={mean_abs:.2e} (~0) — near-zero/unserialized" + ) + + # fp16 concat is lossless, so the strongest assert also holds end-to-end. + np.testing.assert_array_equal( + packed, + expected, + err_msg="packed-QKV bytes corrupted by the serialization round-trip", + ) + + # Functional check: ORT inference on the reloaded model must match a numpy + # reference, proving the packed weight is correct in use, not just on disk. + sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + feed = {"hidden": rng.standard_normal((2, 16)).astype(np.float16)} + (result,) = sess.run(None, feed) + reference = feed["hidden"].astype(np.float32) @ expected.astype(np.float32) + np.testing.assert_allclose(result.astype(np.float32), reference, rtol=1e-2, atol=1e-2) + + def test_value_gate_catches_corrupted_packed_slice(self, tmp_path): + """Negative control: the per-slice value gate must FAIL on a bad pack. + + Proves the discriminators in + ``test_folded_fp16_packed_qkv_values_survive_serialization`` actually + catch the original failure signature — a packed-QKV with a near-zero + slice (corr≈0 / norm collapse) — and that the corruption survives the + serialize→reload round-trip rather than being silently "healed" by + save/load. Without this negative control a future change could neuter + the value asserts and still go green. + """ + rng = np.random.default_rng(1) + q_arr = rng.standard_normal((16, 16)).astype(np.float16) + k_arr = rng.standard_normal((16, 16)).astype(np.float16) + v_arr = rng.standard_normal((16, 16)).astype(np.float16) + + # Poison the K slice to zero — the exact "unserialized / near-zero packed + # weight" signature of the original garbage export. + poisoned = np.concatenate([q_arr, k_arr, v_arr], axis=1).copy() + poisoned[:, 16:32] = 0 + + packed = ir.Value( + name="packed_qkv", + shape=ir.Shape([16, 48]), + type=ir.TensorType(ir.DataType.FLOAT16), + ) + packed.const_value = ir.tensor(poisoned) + graph = ir.Graph( + inputs=[], + outputs=[packed], + nodes=[], + name="poisoned_pack", + opset_imports={"": 20}, + ) + graph.register_initializer(packed) + model = ir.Model(graph, ir_version=10) + + model_path = tmp_path / "model.onnx" + ir.save(model, model_path, external_data="model.onnx.data") + reloaded = ir.load(model_path) + got = reloaded.graph.initializers["packed_qkv"].const_value.numpy() + + # The corruption must survive the round-trip (save/load must not "fix" it). + np.testing.assert_array_equal(got[:, 16:32], np.zeros((16, 16), dtype=np.float16)) + + # The gate's discriminators must flag the zeroed K slice. mean|abs| is the + # robust primary signal: corr is undefined (nan) for a zero-variance slice, + # which is exactly why a corr-only gate would be unsafe here. + k_block = got[:, 16:32].astype(np.float32).ravel() + k_ref = k_arr.astype(np.float32).ravel() + assert np.abs(k_block).mean() < 1e-3, ( + "mean|abs| discriminator must flag a zeroed packed slice" + ) + rel_err = abs(np.linalg.norm(k_block) - np.linalg.norm(k_ref)) / np.linalg.norm(k_ref) + assert rel_err > 0.02, "norm rel_err discriminator must flag a zeroed packed slice" + + # The untouched Q and V slices must still read as healthy — the gate is + # specific to the corrupted slice, not a blanket failure. + for sl, ref in ((got[:, 0:16], q_arr), (got[:, 32:48], v_arr)): + sl32 = sl.astype(np.float32).ravel() + assert np.abs(sl32).mean() >= 1e-3 + corr = np.corrcoef(sl32, ref.astype(np.float32).ravel())[0, 1] + assert corr >= 0.99 diff --git a/src/mobius/_passes/_fold_dtype_e2e_test.py b/src/mobius/_passes/_fold_dtype_e2e_test.py new file mode 100644 index 00000000..0b98f59c --- /dev/null +++ b/src/mobius/_passes/_fold_dtype_e2e_test.py @@ -0,0 +1,358 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""End-to-end regression tests for fp16 dtype preservation through the fold passes. + +These complement the unit-level coverage in ``_dtype_utils_test.py`` and the +pass-level coverage in ``_fold_concat_test.py`` / ``_fold_transpose_test.py`` by +driving the *real* export pipeline (``build_from_module`` + ``apply_weights``) +rather than hand-built single-pass graphs. + +The df203cc fix has two parts, and this module guards both at the export level: + +1. **dtype stamping (type-stamping path).** The fold passes stamp the *resolved* + fp16 dtype on the new packed / transposed initializer (its ``TensorType`` and + its ``LazyTensor``) instead of inheriting the unset type of the rewrite-produced + ``Concat`` intermediate or defaulting to ``FLOAT``. Guarded by + ``test_packed_qkv_weights_stay_fp16_after_export`` and + ``test_all_folded_matmul_weights_stay_fp16_after_export``: in the real fp16 GQA + export the PackQKV rewrite emits a ``MatMul(hidden, Transpose(Concat(W_q, W_k, + W_v)))`` whose ``Concat`` output carries no declared dtype, so without the + stamping the chained FoldConcat -> FoldTranspose widens the packed-QKV weight + to fp32. + +2. **const_value fallback.** ``_dtype_utils.initializer_dtype`` resolves the dtype + from ``const_value`` when an initializer's declared ``type`` was dropped during + graph building. Guarded by + ``test_dropped_declared_dtype_falls_back_to_fp16_const``, which reproduces that + dropped-dtype condition on the packed-QKV weights (as the original regression + exhibited) so the const_value fallback is the only thing keeping the folded + weights fp16. + +Either failure silently widens an fp16 model to fp32 packed weights, which +onnxruntime rejects at load time with a fp16/fp32 ``MatMul`` type-mismatch +(``Type parameter (T) of Optype (MatMul) bound to different types``). + +The in-memory dtype assertions read the IR ``const_value`` metadata. The +regression's *ground-truth* symptom, however, is a serialize-time corruption: +the fold passes write the fp16 *bytes* under a fp32 declaration, so the persisted +initializer is a fp32-declared tensor backed by half-width data. An in-memory +``const_value.numpy()`` can stay fp16 and miss this (only the declared dtype +widens), so ``test_packed_qkv_survives_serialization_roundtrip`` and the +serialize check in ``test_dropped_declared_dtype_falls_back_to_fp16_const`` drive +the real export artifact (``ir.save`` with external data — the same +``model.onnx`` + ``model.onnx.data`` layout the fp16 export produces), reload it, +and assert the folded weights reload as fp16 with their bytes intact. + +The models are tiny and fully synthetic (no HuggingFace download, no GPU and no +onnxruntime execution — only the ONNX graph is constructed, serialized and +inspected), so the tests fit the per-PR CI tier. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest +import torch + +from mobius._builder import build_from_module +from mobius._configs import ArchitectureConfig +from mobius._optimizations import fold_initializers_after_weights +from mobius._registry import registry +from mobius._weight_loading import apply_weights + + +def _make_fp16_llama_config() -> ArchitectureConfig: + """A tiny llama config whose Q/K/V weights are packable (no QK norm). + + ``dtype=FLOAT16`` routes the build through the fp16 packed-QKV path that the + fold passes must preserve. + """ + return ArchitectureConfig( + # Structural invariants that make the fp16 PackQKV fold fire: + # num_attention_heads * head_dim == hidden_size (4 * 16 == 64), and + # num_key_value_heads < num_attention_heads (2 < 4, i.e. GQA). + # The remaining fields are arbitrary small values kept lightweight. + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_hidden_layers=2, + vocab_size=256, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + pad_token_id=0, + dtype=ir.DataType.FLOAT16, + ) + + +def _build_fp16_decoder(config: ArchitectureConfig) -> ir.Model: + """Build an fp16 llama decoder graph with packed QKV, weights not yet loaded. + + Uses the ``cuda`` execution provider so the GQA + PackQKV rewrites pack + Q/K/V into a single ``MatMul(hidden, Transpose(Concat(...)))``. Only the ONNX + graph is constructed — no GPU or onnxruntime execution is required. The fold + passes have not run yet (they run when weights are loaded). + """ + return build_from_module( + registry.get("llama")(config), + config, + execution_provider="cuda", + )["model"] + + +def _fp16_weight_tensors(model: ir.Model) -> dict[str, torch.Tensor]: + """fp16 tensors for every uninitialised parameter, mirroring an fp16 checkpoint. + + Weight initializers always carry fully-static integer shapes (they are + concrete parameter tensors, never symbolic), so ``int(d)`` is safe here. + """ + return { + name: torch.randn(*[int(d) for d in init.shape]).to(torch.float16) + for name, init in model.graph.initializers.items() + if init.const_value is None + } + + +def _matmul_weight_initializers(model: ir.Model) -> list[ir.Value]: + """Return every ``MatMul`` second input that is a graph initializer.""" + initializers = model.graph.initializers + weights: list[ir.Value] = [] + for node in model.graph: + if node.op_type != "MatMul": + continue + weight = node.inputs[1] + if weight is not None and weight.name in initializers: + weights.append(weight) + return weights + + +def _assert_matmul_weights_roundtrip_as_fp16(model: ir.Model, tmp_path: Path) -> None: + """Assert every folded MatMul weight reloads as fp16 with its bytes intact. + + Serializes ``model`` through the production external-data path. This is the + ground-truth df203cc guard. The regression writes the packed / transposed fp16 *bytes* under a fp32 declaration, a corruption an in-memory + ``const_value`` check can mask (the backing array stays fp16 while only the + declared dtype widens) but that surfaces on serialize -> reload as a + fp32-declared initializer whose bytes no longer match the fp16 source. + ``ir.save`` with ``external_data`` mirrors the real fp16 export's + ``model.onnx`` + ``model.onnx.data`` layout, exactly where the dtype bug + corrupts the weight bytes. + """ + before = {w.name: w.const_value.numpy() for w in _matmul_weight_initializers(model)} + assert before, "Expected the export to contain MatMul weight initializers to round-trip" + + model_path = tmp_path / "model.onnx" + ir.save(model, model_path, external_data="model.onnx.data") + assert (tmp_path / "model.onnx.data").exists(), ( + "fp16 weights should be externalized, exercising the real export's " + "model.onnx + model.onnx.data save path where the dtype bug corrupts bytes" + ) + + reloaded = ir.load(model_path) + for weight in _matmul_weight_initializers(reloaded): + assert weight.dtype == ir.DataType.FLOAT16, ( + f"Reloaded MatMul weight {weight.name!r} serialized as {weight.dtype} " + f"(expected FLOAT16): fp16 bytes were written under a fp32 declaration " + f"— the df203cc serialize-time corruption." + ) + np.testing.assert_array_equal( + weight.const_value.numpy(), + before[weight.name], + err_msg=( + f"Reloaded MatMul weight {weight.name!r} bytes changed across " + f"serialize -> reload; the fp16 packed weight was silently corrupted " + f"(fp16 data written under a fp32 dtype)." + ), + ) + + +@pytest.fixture +def fp16_export() -> tuple[ArchitectureConfig, ir.Model]: + """A real fp16 packed-QKV export: build + ``apply_weights`` (folds inside). + + Function-scoped (a fresh build per test) on purpose: the serialize-roundtrip + test calls ``ir.save`` on the model, and on some ``onnx_ir`` versions that + offloads the initializers' ``const_value`` to external tensors in place. A + shared (module-scoped) model could then leak that mutated/externalized state + into another test, and under ``pytest-xdist`` the tests' execution order is + not guaranteed — making such cross-test contamination order-dependent and + flaky. A fresh model per test keeps each test hermetic and deterministic. The + build is tiny, so the rebuild cost is negligible. + """ + config = _make_fp16_llama_config() + model = _build_fp16_decoder(config) + # apply_weights assigns the fp16 weights and then runs + # fold_initializers_after_weights, folding the Transpose/Concat weight nodes. + apply_weights(model, _fp16_weight_tensors(model)) + return config, model + + +class TestFp16FoldDtypeE2E: + def test_packed_qkv_weights_stay_fp16_after_export( + self, fp16_export: tuple[ArchitectureConfig, ir.Model] + ) -> None: + """An fp16 packed-QKV export keeps fp16 packed/transposed weights. + + Guards the df203cc dtype-stamping path: the chained + ``FoldConcatInitializersPass`` (the Q/K/V pack) and + ``FoldTransposedInitializerPass`` (the weight transpose) must carry the + fp16 dtype through the rewrite-produced ``Concat`` intermediate instead + of defaulting it to ``FLOAT``. + """ + config, model = fp16_export + + # The packed path must actually have run, otherwise the dtype assertion + # below would pass vacuously (e.g. if the build stopped packing QKV). + gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) == config.num_hidden_layers, ( + "Expected one packed GroupQueryAttention per layer; the fp16 GQA " + "PackQKV path was not exercised, so this regression guard is vacuous." + ) + + for gqa in gqa_nodes: + matmul = gqa.inputs[0].producer() + assert matmul is not None and matmul.op_type == "MatMul", ( + "Packed GQA projection should be produced by a MatMul" + ) + weight = matmul.inputs[1] + assert weight is not None, "Packed projection MatMul is missing its weight input" + # FoldConcat + FoldTranspose ran inside apply_weights, so the packed + # weight is now a plain folded initializer (no producer node left). + assert weight.producer() is None, ( + "Packed QKV weight should be a folded initializer after export, " + "not the output of a residual Transpose/Concat node" + ) + assert weight.const_value is not None + assert weight.const_value.dtype == ir.DataType.FLOAT16, ( + f"Packed QKV weight {weight.name!r} widened to " + f"{weight.const_value.dtype} (expected FLOAT16); the fp16 fold " + f"dtype fix (df203cc) has regressed." + ) + assert weight.dtype == ir.DataType.FLOAT16 + + def test_all_folded_matmul_weights_stay_fp16_after_export( + self, fp16_export: tuple[ArchitectureConfig, ir.Model] + ) -> None: + """Every folded MatMul weight in an fp16 export stays fp16. + + Broader guard than the packed-QKV case: it also covers the standalone + ``o_proj`` and MLP weight transposes folded by + ``FoldTransposedInitializerPass``. A single fp32-widened weight here is + exactly what makes onnxruntime reject the fp16 model at load time. + """ + _, model = fp16_export + + weights = _matmul_weight_initializers(model) + assert weights, "Expected the export to contain MatMul weight initializers" + # Anti-vacuity: at least one weight must be a folded transposed initializer + # (FoldTransposedInitializerPass names them ``..._t``). Without this, the + # test could pass simply because no folding ran. + assert any(weight.name.endswith("_t") for weight in weights), ( + "Expected at least one folded transposed weight (``..._t``); the fold " + "passes did not run, so this regression guard would be vacuous." + ) + + # fp32 widening (FLOAT) is the specific documented failure mode of the + # df203cc regression — the fold passes defaulting a dropped declared dtype + # to FLOAT — so we flag exactly that rather than any non-fp16 dtype. + widened = [ + weight.name + for weight in weights + if weight.const_value is not None and weight.const_value.dtype == ir.DataType.FLOAT + ] + assert not widened, ( + f"fp16 export produced fp32 MatMul weights after folding: {widened}. " + f"The fold passes must preserve fp16 (df203cc)." + ) + + def test_packed_qkv_survives_serialization_roundtrip( + self, fp16_export: tuple[ArchitectureConfig, ir.Model], tmp_path: Path + ) -> None: + """The fp16 export's folded weights survive serialize -> reload as fp16. + + Ground-truth guard for the dtype-stamping path. The in-memory assertions + above read the IR ``const_value`` metadata; this drives the real export + artifact (``ir.save`` with external data, like the production + ``model.onnx`` + ``model.onnx.data``) and reloads it, the only check that + catches the df203cc serialize-time fp16-under-fp32 byte corruption — an + in-memory ``const_value.numpy()`` can stay fp16 and miss it. + """ + _, model = fp16_export + _assert_matmul_weights_roundtrip_as_fp16(model, tmp_path) + + def test_dropped_declared_dtype_falls_back_to_fp16_const(self, tmp_path: Path) -> None: + """Folding fp16 weights whose declared dtype was dropped stays fp16. + + Guards the df203cc ``initializer_dtype`` const_value fallback. The + original regression arose because graph building left the packed-QKV + weights with fp16 ``const_value`` but no declared ``type``; the fold + passes then defaulted to ``FLOAT``. Here we reproduce that exact + condition on a real fp16 export — load fp16 weights, clear the declared + type on the packed-QKV ``Concat`` inputs, then run the real fold + orchestration — so the const_value fallback is the *only* thing that can + keep the folded weights fp16. + + ``fold_initializers_after_weights`` is invoked directly (rather than via + ``apply_weights``) so the dropped-dtype condition can be injected between + loading the weights and folding; it is the same orchestration + ``apply_weights`` runs internally. + """ + config = _make_fp16_llama_config() + model = _build_fp16_decoder(config) + + # Load fp16 weights onto the initializers (mirrors apply_weights' assign + # step) without yet folding. + for name, tensor in _fp16_weight_tensors(model).items(): + model.graph.initializers[name].const_value = ir.tensor(tensor.numpy()) + + # Reproduce the dropped-declared-dtype condition: clear the declared type + # on the packed-QKV weight initializers (the Concat inputs) while keeping + # their fp16 const_value. + dropped: list[str] = [] + for node in model.graph: + if node.op_type != "Concat": + continue + for inp in node.inputs: + if ( + inp is not None + and inp.name in model.graph.initializers + and inp.const_value is not None + ): + inp.type = None + dropped.append(inp.name) + assert dropped, ( + "Expected packed-QKV Concat inputs to drop declared dtype on; the " + "fp16 PackQKV path was not exercised, so this guard would be vacuous." + ) + + fold_initializers_after_weights(model) + + gqa_nodes = [n for n in model.graph if n.op_type == "GroupQueryAttention"] + assert len(gqa_nodes) == config.num_hidden_layers + for gqa in gqa_nodes: + matmul = gqa.inputs[0].producer() + assert matmul is not None and matmul.op_type == "MatMul", ( + "Packed GQA projection should be produced by a MatMul" + ) + weight = matmul.inputs[1] + assert weight is not None, "Packed projection MatMul is missing its weight input" + assert weight.const_value is not None + assert weight.const_value.dtype == ir.DataType.FLOAT16, ( + f"Packed QKV weight {weight.name!r} widened to " + f"{weight.const_value.dtype} when its declared dtype was dropped; " + f"the initializer_dtype const_value fallback (df203cc) regressed." + ) + + # Ground-truth check: the const_value fallback must also hold through the + # real serialize -> reload path. Reverting only the ``initializer_dtype`` + # call-sites widens these dropped-dtype weights and corrupts their bytes + # here even though the type-stamp keeps the realistic-export tests green. + _assert_matmul_weights_roundtrip_as_fp16(model, tmp_path) diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py index 402cad7c..e7f6cae3 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -26,6 +26,8 @@ import onnx_ir as ir +from mobius._passes._dtype_utils import initializer_dtype + logger = logging.getLogger(__name__) @@ -107,6 +109,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: new_name = f"{inp.name}_t" + # Resolve the dtype from the declared type or, when it was + # dropped during graph building, from the loaded const_value, so + # a transposed fp16 weight is not silently emitted as fp32. + transposed_dtype = initializer_dtype(inp) or ir.DataType.FLOAT + # Derive shape of the transposed tensor from the Transpose output. # Fall back to computing from the input shape if shape inference # did not propagate to this new node (e.g. after stage-2 rewrites). @@ -115,7 +122,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: perm = list(perm_attr.value) t_shape = ir.Shape([inp.shape[p] for p in perm]) - new_val = ir.Value(name=new_name, shape=t_shape, type=inp.type) + new_val = ir.Value( + name=new_name, + shape=t_shape, + type=ir.TensorType(transposed_dtype), + ) # Create a LazyTensor that transposes the original data on demand. # The actual numpy transposition is deferred until serialization, @@ -123,7 +134,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: src = inp # captured for the closure below new_val.const_value = ir.LazyTensor( lambda s=src: ir.tensor(s.const_value.numpy().T), # type: ignore[union-attr] - dtype=inp.dtype or ir.DataType.FLOAT, + dtype=transposed_dtype, shape=t_shape, name=new_val.name, ) diff --git a/src/mobius/_passes/_fold_transpose_test.py b/src/mobius/_passes/_fold_transpose_test.py index 93c74b37..0e3def2a 100644 --- a/src/mobius/_passes/_fold_transpose_test.py +++ b/src/mobius/_passes/_fold_transpose_test.py @@ -508,3 +508,53 @@ def test_original_weight_pruned_after_full_pipeline(): assert "weight_t" not in model.graph.initializers transpose_nodes = [n for n in model.graph.all_nodes() if n.op_type == "Transpose"] assert len(transpose_nodes) == 1 + + +class TestFoldTransposeDtype: + """Regression: a transposed fp16 weight with unset declared dtype stays fp16.""" + + def test_transposed_dtype_follows_const_value_when_declared_dtype_missing(self): + """fp16 weights lose their declared dtype after _cast_module_dtype. + + The pass must resolve the transposed initializer's dtype from + ``const_value`` rather than defaulting to FLOAT, otherwise a fp16 weight + is emitted as fp32 and ORT rejects the downstream MatMul. + """ + weight = ir.Value(name="weight") + weight.const_value = ir.tensor(np.arange(12, dtype=np.float16).reshape(4, 3)) + # Do NOT set .dtype — declared dtype stays None. + assert weight.dtype is None + + x = ir.Value(name="x") + x.shape = ir.Shape([2, 4]) + x.dtype = ir.DataType.FLOAT16 + + transpose_node = ir.Node( + "", + "Transpose", + inputs=[weight], + attributes=[ir.Attr("perm", ir.AttributeType.INTS, [1, 0])], + num_outputs=1, + ) + w_t = transpose_node.outputs[0] + matmul_node = ir.Node("", "MatMul", inputs=[x, w_t], num_outputs=1) + + graph = ir.Graph( + inputs=[x], + outputs=[matmul_node.outputs[0]], + nodes=[transpose_node, matmul_node], + name="test_graph", + opset_imports={"": 20}, + ) + graph.register_initializer(weight) + model = ir.Model(graph, ir_version=10) + + result = FoldTransposedInitializerPass()(model) + assert result.modified + + packed = model.graph.initializers["weight_t"] + assert packed.dtype == ir.DataType.FLOAT16, ( + f"Transposed initializer declared dtype should be FLOAT16, got {packed.dtype}" + ) + assert packed.const_value.dtype == ir.DataType.FLOAT16 + assert packed.const_value.numpy().dtype == np.float16 diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index 41cf5071..d8c156b5 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -114,13 +114,69 @@ def _register_kv_cache_outputs( present_key_values: list[tuple[ir.Value, ir.Value]], *, prefix: str = "present", + batch: ir.SymbolicDim | str | None = None, + num_kv_heads: int | None = None, + key_head_dim: int | None = None, + value_head_dim: int | None = None, + total_seq_len: ir.SymbolicDim | str | int | None = None, + dtype: ir.DataType | None = None, ) -> None: """Name and register KV cache outputs on the graph. - Output shapes and dtypes are inferred by the shape inference pass - that runs during model optimization. + When every present-shape parameter (``batch``, ``num_kv_heads``, + ``key_head_dim``, ``value_head_dim``, ``total_seq_len``, ``dtype``) is + supplied, each ``present.{i}.{key,value}`` output is stamped with an + explicit ``[batch, num_kv_heads, total_seq_len, head_dim]`` type, symmetric + to the ``past_key_values`` inputs created by :func:`_make_kv_cache_inputs`. + + This explicit stamp is required for ``com.microsoft::GroupQueryAttention`` + exports: that contrib op's shape inference mis-derives the present + ``head_dim`` (it divides the *packed* QKV query hidden by + ``num_heads + 2 * kv_num_heads`` and lands on the wrong value), so the + ``present.*`` outputs would otherwise declare the wrong ``head_dim`` even + though the kernel produces correct data at runtime. The mismatch makes ORT + log a shape-merge warning and breaks present->past chaining in consumers + such as onnxruntime-genai that trust the declared shapes. The plain ONNX + ``Attention`` op infers the present shape correctly, so the stamp is a + no-op there. + + When the parameters are omitted, output shapes/dtypes are left to the + shape inference pass that runs during model optimization. + + The present-shape parameters are all-or-nothing by design: pass every one + to stamp the explicit type, or none to opt out and infer. A *partial* set + is always a wiring slip (a caller wired some dims but dropped others) and is + never legitimate, so it is rejected fail-closed: a partial set raises + :class:`ValueError` naming the provided and missing parameters, rather than + silently falling back to the known-wrong inference path. """ + params = { + "batch": batch, + "num_kv_heads": num_kv_heads, + "key_head_dim": key_head_dim, + "value_head_dim": value_head_dim, + "total_seq_len": total_seq_len, + "dtype": dtype, + } + provided = [name for name, value in params.items() if value is not None] + stamp = len(provided) == len(params) + if provided and not stamp: + missing = [name for name in params if params[name] is None] + raise ValueError( + f"_register_kv_cache_outputs received a partial set of present-shape " + f"parameters (provided {provided}, missing {missing}); these are " + f"all-or-nothing. Pass all six to stamp explicit present.* types " + f"(required for correct GroupQueryAttention head_dim), or none to opt " + f"out and infer." + ) for i, (present_key, present_value) in enumerate(present_key_values): + if stamp: + present_key.shape = ir.Shape([batch, num_kv_heads, total_seq_len, key_head_dim]) + present_key.type = ir.TensorType(dtype) + present_value.shape = ir.Shape( + [batch, num_kv_heads, total_seq_len, value_head_dim] + ) + present_value.type = ir.TensorType(dtype) builder.add_output(present_key, f"{prefix}.{i}.key") builder.add_output(present_value, f"{prefix}.{i}.value") diff --git a/src/mobius/tasks/_cache_utils_test.py b/src/mobius/tasks/_cache_utils_test.py new file mode 100644 index 00000000..780e1cfa --- /dev/null +++ b/src/mobius/tasks/_cache_utils_test.py @@ -0,0 +1,124 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for KV-cache I/O helpers in :mod:`mobius.tasks._cache_utils`.""" + +from __future__ import annotations + +import onnx_ir as ir +import pytest + +from mobius.tasks._base import _make_graph +from mobius.tasks._cache_utils import _register_kv_cache_outputs + + +def _present_pair(name: str, wrong_head_dim: int) -> tuple[ir.Value, ir.Value]: + """A present key/value pair carrying a deliberately wrong inferred shape. + + Mimics ``com.microsoft::GroupQueryAttention`` shape inference, which + mis-derives ``head_dim`` (e.g. 32 instead of 96) for the present outputs. + """ + key = ir.Value(name=f"{name}_key") + value = ir.Value(name=f"{name}_value") + for v in (key, value): + v.shape = ir.Shape(["batch", 32, "seq", wrong_head_dim]) + v.type = ir.TensorType(ir.DataType.FLOAT16) + return key, value + + +def _dims(value: ir.Value) -> list[object]: + return [d if isinstance(d, int) else str(d) for d in value.shape] + + +class TestRegisterKVCacheOutputs: + def test_stamps_explicit_shapes_over_wrong_inference(self): + """Explicit params must override mis-inferred present shapes (GQA bug).""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + _register_kv_cache_outputs( + builder, + pairs, + batch=ir.SymbolicDim("batch"), + num_kv_heads=32, + key_head_dim=96, + value_head_dim=96, + total_seq_len="past_sequence_len + sequence_len", + dtype=ir.DataType.FLOAT16, + ) + + key, value = pairs[0] + assert _dims(key) == ["batch", 32, "past_sequence_len + sequence_len", 96] + assert _dims(value) == ["batch", 32, "past_sequence_len + sequence_len", 96] + assert key.dtype == ir.DataType.FLOAT16 + assert value.dtype == ir.DataType.FLOAT16 + + def test_distinct_key_value_head_dims(self): + """MLA-style caches may use different key/value head dims.""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + _register_kv_cache_outputs( + builder, + pairs, + batch=ir.SymbolicDim("batch"), + num_kv_heads=16, + key_head_dim=192, + value_head_dim=128, + total_seq_len="past_sequence_len + sequence_len", + dtype=ir.DataType.FLOAT16, + ) + + key, value = pairs[0] + assert _dims(key)[1] == 16 and _dims(key)[3] == 192 + assert _dims(value)[1] == 16 and _dims(value)[3] == 128 + + def test_no_params_leaves_shapes_untouched(self): + """Without shape params the helper must not stamp (inference path).""" + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + _register_kv_cache_outputs(builder, pairs) + + key, _ = pairs[0] + # Unchanged: still the (wrong) pre-existing inferred shape. + assert _dims(key) == ["batch", 32, "seq", 32] + + def test_partial_params_raise(self): + """A partial present-shape set is a wiring slip -> must raise. + + This is the fail-closed regression proof: the exact input that + previously warned-and-proceeded (shipping a structurally-wrong model + with mis-derived GroupQueryAttention ``head_dim``) now raises before + any output is registered. + """ + _, builder = _make_graph() + pairs = [_present_pair("present.0", wrong_head_dim=32)] + + with pytest.raises(ValueError, match="partial set of present-shape parameters") as exc: + _register_kv_cache_outputs( + builder, + pairs, + batch=ir.SymbolicDim("batch"), + num_kv_heads=32, + # key_head_dim / value_head_dim / total_seq_len / dtype omitted + ) + + # The message must name every omitted parameter so the slip is diagnosable. + for missing in ("key_head_dim", "value_head_dim", "total_seq_len", "dtype"): + assert missing in str(exc.value) + + def test_registers_named_outputs(self): + """Outputs are registered with the conventional present.{i}.* names.""" + _, builder = _make_graph() + pairs = [_present_pair("a", 32), _present_pair("b", 32)] + + _register_kv_cache_outputs(builder, pairs) + + names = [v.name for v in builder.graph.outputs] + assert names == [ + "present.0.key", + "present.0.value", + "present.1.key", + "present.1.value", + ] diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py index 471d350b..972437c7 100644 --- a/src/mobius/tasks/_causal_lm.py +++ b/src/mobius/tasks/_causal_lm.py @@ -153,6 +153,10 @@ def build( num_kv_cache_heads = ( config.num_attention_heads if use_mla else config.num_key_value_heads ) + kv_key_head_dim = ( + (config.qk_nope_head_dim or 0) + (config.qk_rope_head_dim or 0) + ) or config.head_dim + kv_value_head_dim = config.v_head_dim or config.head_dim past_key_values = _make_kv_cache_inputs( builder, @@ -162,9 +166,8 @@ def build( config.dtype, batch, past_seq_len, - key_head_dim=((config.qk_nope_head_dim or 0) + (config.qk_rope_head_dim or 0)) - or None, - value_head_dim=config.v_head_dim or None, + key_head_dim=kv_key_head_dim, + value_head_dim=kv_value_head_dim, ) logits, present_key_values = module( @@ -184,9 +187,19 @@ def build( present_key_values, ) else: + # Stamp explicit present shapes symmetric to the past inputs so the + # com.microsoft::GroupQueryAttention export declares the correct + # head_dim (its contrib-op shape inference otherwise mis-derives it + # from the packed QKV hidden). total_seq = past + current sequence. _register_kv_cache_outputs( builder, present_key_values, + batch=batch, + num_kv_heads=num_kv_cache_heads, + key_head_dim=kv_key_head_dim, + value_head_dim=kv_value_head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, ) return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/src/mobius/tasks/_task_test.py b/src/mobius/tasks/_task_test.py index d50f4337..4a638a9c 100644 --- a/src/mobius/tasks/_task_test.py +++ b/src/mobius/tasks/_task_test.py @@ -91,6 +91,38 @@ def test_build_outputs(self): assert "present.0.key" in output_names assert "present.0.value" in output_names + def test_present_outputs_match_past_inputs(self): + """present.{i}.* must match the corresponding past_key_values.{i}.* inputs. + + They must declare the same kv_heads/head_dim/dtype as the corresponding + past_key_values.{i}.* inputs, with total (past+current) sequence length. + Guards the GQA present-head_dim export bug, where the contrib-op shape + inference would otherwise mis-declare head_dim. + """ + config = make_config() + module = CausalLMModel(config) + pkg = CausalLMTask().build(module, config) + model = pkg["model"] + inputs = {v.name: v for v in model.graph.inputs} + outputs = {v.name: v for v in model.graph.outputs} + + def dims(value): + return [d if isinstance(d, int) else str(d) for d in value.shape] + + for i in range(config.num_hidden_layers): + for kind in ("key", "value"): + past = inputs[f"past_key_values.{i}.{kind}"] + present = outputs[f"present.{i}.{kind}"] + past_dims = dims(past) + present_dims = dims(present) + # batch, kv_heads, head_dim must match the past input exactly. + assert present_dims[0] == past_dims[0] + assert present_dims[1] == past_dims[1] + assert present_dims[3] == past_dims[3] + # present covers past + current tokens. + assert present_dims[2] == "past_sequence_len + sequence_len" + assert present.dtype == past.dtype + def test_build_producer_info(self): config = make_config() module = CausalLMModel(config)