Reuse quantized embedding table for tied LM head in TieWordEmbeddings

justinchuby · Copilot · justinchuby · commit 5c8867b43900 · 2026-07-01T18:01:34.000Z
Add a reuse mode to the TieWordEmbeddings graph surgery for the case where the
embedding has been quantized to GatherBlockQuantized but the LM head is still a
float MatMul (its weight is the tied embedding, reached through a Transpose).

Previously TieWordEmbeddings only handled both-unquantized (Gather + MatMul) or
both-quantized (GatherBlockQuantized + MatMulNBits). When only the embedding
Gather is quantized (e.g. OnnxBlockWiseRtnQuantization while the body is left for
OnnxKQuantQuantization), the tied word-embedding matrix ends up stored twice: once
as INT4 (embedding) and once as float16 (LM head), which is larger than a fully
float16 model.

handle_reuse rebuilds the LM head as a MatMulNBits that shares the embedding's
INT4 qweight / scales / zero-point (the byte-identical table, reshaped to the
MatMulNBits layout), and prunes the now-dead Transpose and float embedding weight.
reuse_weights_match gates this on the float LM head weight actually matching the
dequantized embedding table, so an untied projection is never tied.

This lets a K-Quant body + shared-INT4 tied embedding/LM head model reach the
smallest on-disk size at the highest-quality body quantization.

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
Signed-off-by: Justin Chu &lt;11205048+justinchuby@users.noreply.github.com&gt;
diff --git a/olive/passes/onnx/graph_surgeries.py b/olive/passes/onnx/graph_surgeries.py
@@ -2128,8 +2128,17 @@ def __call__(self, model: ModelProto):
 class TieWordEmbeddings(ProtoSurgeon):
     """Tie word embeddings.
 
-    Only supports when the embeddings are not quantized or when both are quantized
-    using GatherBlockQuantized and MatMulNBits
+    Supports three cases:
+
+    * both unquantized (``Gather`` + ``MatMul``),
+    * both quantized (``GatherBlockQuantized`` + ``MatMulNBits``), and
+    * a *reuse* case where the embedding is quantized (``GatherBlockQuantized``)
+      but the LM head is still a float ``MatMul``. The LM head is rebuilt as a
+      ``MatMulNBits`` that shares the embedding's quantized weight table, so the
+      large word-embedding matrix is stored once as INT4 instead of once as INT4
+      (embedding) plus once as float16 (LM head). This is useful after quantizing
+      only the embedding ``Gather`` (e.g. ``OnnxBlockWiseRtnQuantization`` with the
+      body left for a separate pass such as ``OnnxKQuantQuantization``).
     """
 
     def __call__(self, model: onnx.ModelProto):
@@ -2151,9 +2160,19 @@ def __call__(self, model: onnx.ModelProto):
         if embed_name is None:
             return dag.model
 
-        lm_head_name, lm_head_op_type = self.get_name_op_type(
-            dag, [dag.get_producer("logits")], ["MatMul", "MatMulNBits"], 1
-        )
+        # Reuse case: embedding already quantized but the LM head is still a float
+        # MatMul. Its weight is the tied embedding, possibly reached through a
+        # Transpose, so the weight input is not necessarily a direct initializer
+        # (which get_name_op_type would require). Detect it directly here.
+        logits_producer = dag.get_producer("logits")
+        if (
+            embed_op_type == "GatherBlockQuantized"
+            and logits_producer is not None
+            and dag.get_node_op_type(logits_producer) == "MatMul"
+        ):
+            return self.handle_reuse(dag, embed_name, logits_producer)
+
+        lm_head_name, lm_head_op_type = self.get_name_op_type(dag, [logits_producer], ["MatMul", "MatMulNBits"], 1)
         if lm_head_name is None:
             return dag.model
 
@@ -2348,6 +2367,176 @@ def handle_quantized(self, dag: OnnxDAG, embed_name: str, lm_head_name: str):
         dag.update()
         return dag.model
 
+    def handle_reuse(self, dag: OnnxDAG, embed_name: str, lm_head_name: str):
+        """Rebuild a float ``MatMul`` LM head as a ``MatMulNBits`` reusing the embedding's INT4 table.
+
+        The embedding is a ``GatherBlockQuantized`` node with inputs
+        ``[qweight, indices, scales, zero_point?]``; the LM head is a float
+        ``MatMul(activation, W)`` whose weight ``W`` is the tied embedding matrix,
+        possibly reached through a ``Transpose``. Because ``W`` equals the
+        embedding table, its quantized form is exactly the embedding's INT4 table,
+        so we point the new ``MatMulNBits`` at the embedding's ``qweight`` / scales
+        / zero-point instead of quantizing ``W`` a second time. The float ``W``
+        (and any feeding ``Transpose``) then becomes dead and is pruned.
+        """
+        embed_inputs = dag.get_node_inputs(embed_name)
+        embed_qweight = embed_inputs[0]
+        embed_scales = embed_inputs[2]
+        embed_zero_point = embed_inputs[3] if len(embed_inputs) > 3 else None
+
+        embed_attrs = dag.get_node_attributes(embed_name)
+        bits = int(embed_attrs["bits"])
+        block_size = int(embed_attrs["block_size"])
+
+        # LM head weight: trace through an optional Transpose to the source initializer.
+        lm_head_inputs = dag.get_node_inputs(lm_head_name)
+        lm_head_activation = lm_head_inputs[0]
+        weight_input = lm_head_inputs[1]
+        fp16_weight = weight_input
+        transpose_name = None
+        if not dag.is_initializer(weight_input):
+            weight_producer = dag.get_producer(weight_input)
+            if weight_producer is not None and dag.get_node_op_type(weight_producer) == "Transpose":
+                transpose_name = weight_producer
+                fp16_weight = dag.get_node_inputs(weight_producer)[0]
+        if not dag.is_initializer(fp16_weight):
+            logger.debug("LM head weight is not a (transposed) initializer; cannot reuse embedding quantization.")
+            return dag.model
+
+        # embed qweight is [vocab, hidden * bits / 8]; derive N (vocab) and K (hidden).
+        embed_qweight_shape = dag.get_io_shape(embed_qweight)
+        vocab = int(embed_qweight_shape[0])
+        hidden = int(embed_qweight_shape[1]) * 8 // bits
+
+        # Correctness gate: only reuse when the float LM head weight really is the
+        # (dequantized) embedding table, so an untied projection is never tied.
+        if not self.reuse_weights_match(dag, embed_name, fp16_weight, vocab, hidden, bits, block_size):
+            logger.debug("LM head weight does not match the embedding table; skipping reuse tie.")
+            return dag.model
+
+        logger.debug("Reusing quantized embedding weights for the LM head")
+
+        graph_idx = dag.get_graph_idx(lm_head_name)
+        n_blocks = hidden // block_size
+        blob_size = block_size * bits // 8
+
+        # GatherBlockQuantized stores qweight 2D [vocab, hidden * bits / 8]; MatMulNBits
+        # expects 3D [N, n_blocks, blob_size]. The bytes are identical, only reshaped.
+        reshape_output = self.add_reshape_node(
+            dag,
+            graph_idx,
+            self.create_new_name(lm_head_name, "MatMul", "Reshape_tied"),
+            embed_qweight,
+            [vocab, n_blocks, blob_size],
+            dag.get_io_elem_type(embed_qweight),
+        )
+
+        mmnb_inputs = [lm_head_activation, reshape_output, embed_scales]
+        if embed_zero_point is not None:
+            mmnb_inputs.append(embed_zero_point)
+
+        lm_head_output = dag.get_node_outputs(lm_head_name)[0]
+        mmnb_name = self.create_new_name(lm_head_name, "MatMul", "MatMulNBits")
+        mmnb_output = f"{mmnb_name}_output"
+        dag.add_node(
+            onnx.helper.make_node(
+                "MatMulNBits",
+                mmnb_inputs,
+                [mmnb_output],
+                name=mmnb_name,
+                domain="com.microsoft",
+                K=hidden,
+                N=vocab,
+                bits=bits,
+                block_size=block_size,
+            ),
+            graph_idx,
+        )
+
+        # Redirect consumers of the old logits, then swap the graph output over.
+        # remove_node deletes the output IO, so preserve the logits value_info by
+        # copying it onto the new output before renaming (as in handle_unquantized).
+        mmnb_output_vi = onnx.ValueInfoProto()
+        mmnb_output_vi.CopyFrom(dag.get_value_info_proto(lm_head_output))
+        mmnb_output_vi.name = mmnb_output
+        dag.add_value_info(mmnb_output_vi, graph_idx)
+
+        for consumer in dag.get_consumers(lm_head_name):
+            dag.replace_node_input(consumer, lm_head_output, mmnb_output)
+        dag.remove_output(lm_head_output)
+        dag.remove_node(lm_head_name)
+        # Drop the now-orphaned Transpose (if any) so the float embedding weight it
+        # held loses its last consumer and is pruned from the serialized model.
+        if transpose_name is not None and not dag.get_consumers(transpose_name):
+            dag.remove_node(transpose_name)
+        dag.rename_node_output(mmnb_name, mmnb_output, lm_head_output)
+        dag.make_output(lm_head_output)
+
+        dag.update()
+        return dag.model
+
+    def reuse_weights_match(
+        self,
+        dag: OnnxDAG,
+        embed_name: str,
+        fp16_weight_name: str,
+        vocab: int,
+        hidden: int,
+        bits: int,
+        block_size: int,
+        n_check: int = 256,
+    ) -> bool:
+        """Check that ``fp16_weight`` equals the dequantized embedding table.
+
+        Only a slice of ``n_check`` rows is compared, which is enough to reject an
+        untied projection while keeping the check cheap. Only 4-bit packing is
+        supported; other widths conservatively return ``False``.
+        """
+        if bits != 4:
+            return False
+
+        embed_inputs = dag.get_node_inputs(embed_name)
+        qweight = dag.get_initializer_np_array(embed_inputs[0])  # [vocab, hidden * bits / 8] uint8
+        scales = dag.get_initializer_np_array(embed_inputs[2]).astype(np.float32)  # [vocab, n_blocks]
+        zero_point = dag.get_initializer_np_array(embed_inputs[3]) if len(embed_inputs) > 3 else None
+
+        n = min(n_check, vocab)
+        n_blocks = hidden // block_size
+
+        # Unpack two 4-bit codes per byte (low nibble first), matching MatMulNBits packing.
+        q = qweight[:n]
+        codes = np.empty((n, hidden), np.float32)
+        codes[:, 0::2] = (q & 0x0F).astype(np.float32)
+        codes[:, 1::2] = (q >> 4).astype(np.float32)
+        codes = codes.reshape(n, n_blocks, block_size)
+
+        if zero_point is not None:
+            zp = zero_point[:n]
+            zcodes = np.empty((n, n_blocks), np.float32)
+            zcodes[:, 0::2] = (zp & 0x0F).astype(np.float32)
+            zcodes[:, 1::2] = (zp >> 4).astype(np.float32)
+            zcodes = zcodes.reshape(n, n_blocks, 1)
+        else:
+            # Symmetric quantization centers codes at the midpoint of the 4-bit range.
+            zcodes = np.float32(2 ** (bits - 1))
+
+        block_scales = scales[:n].reshape(n, n_blocks, 1)
+        dequant = ((codes - zcodes) * block_scales).reshape(n, hidden)  # ~= embedding[:n]
+
+        fp16 = dag.get_initializer_np_array(fp16_weight_name).astype(np.float32)
+        if list(fp16.shape) == [hidden, vocab]:
+            ref = fp16[:, :n].T
+        elif list(fp16.shape) == [vocab, hidden]:
+            ref = fp16[:n]
+        else:
+            return False
+
+        # Quantization error per element is bounded by one block step; use a
+        # generous multiple so a genuinely tied weight always passes while an
+        # untied projection (differing by ~weight magnitude) is rejected.
+        atol = max(4.0 * float(block_scales.max()), 1e-2)
+        return np.allclose(dequant, ref, atol=atol, rtol=0.0)
+
     def equal_weights(self, dag: OnnxDAG, init0: str, init1: str, transpose: bool = False) -> bool:
         shape0, shape1 = dag.get_io_shape(init0), dag.get_io_shape(init1)
         if np.prod(shape0) != np.prod(shape1):
diff --git a/test/passes/onnx/test_graph_surgeries.py b/test/passes/onnx/test_graph_surgeries.py
@@ -2164,6 +2164,135 @@ def test_tie_word_embeddings(tmp_path, quantized):
         )
 
 
+def _quantize_blockwise_int4(weight, block_size):
+    """Block-wise asymmetric 4-bit RTN quantization matching MatMulNBits packing.
+
+    Quantizes ``weight`` ([rows, cols], float) along ``cols`` in blocks and returns
+    ``(qweight, scales, zero_point)`` with the same packing (two 4-bit codes per
+    byte, low nibble first) that ``TieWordEmbeddings.reuse_weights_match`` expects.
+    """
+    rows, cols = weight.shape
+    n_blocks = cols // block_size
+    w = weight.reshape(rows, n_blocks, block_size)
+    w_min = w.min(axis=2)
+    w_max = w.max(axis=2)
+    scale = np.where((w_max - w_min) == 0, 1.0, (w_max - w_min) / 15.0)
+    zero = np.clip(np.rint(-w_min / scale), 0, 15)
+    codes = np.clip(np.rint(w / scale[:, :, None] + zero[:, :, None]), 0, 15).astype(np.uint8)
+    codes = codes.reshape(rows, cols)
+    # Pack two 4-bit codes per byte (even column -> low nibble).
+    qweight = (codes[:, 0::2] | (codes[:, 1::2] << 4)).astype(np.uint8)  # [rows, cols/2]
+    zcodes = zero.astype(np.uint8)  # [rows, n_blocks]
+    zero_point = (zcodes[:, 0::2] | (zcodes[:, 1::2] << 4)).astype(np.uint8)  # [rows, n_blocks/2]
+    return qweight, scale.astype(np.float16), zero_point
+
+
+def _build_quantized_embedding_tied_model(tmp_path, lm_head_weight):
+    """Build a tiny model with a GatherBlockQuantized embedding and a float MatMul LM head.
+
+    The LM head computes ``hidden @ Transpose(lm_head_weight)`` so its weight is
+    reached through a Transpose, mirroring how mobius emits tied embeddings. The
+    embedding is quantized to INT4; ``lm_head_weight`` ([vocab, hidden], float16) is
+    the tied weight (equal to the embedding table) or an unrelated weight.
+    """
+    vocab, hidden, block_size = 6, 8, 4
+    batch, seq = 1, 2
+
+    qweight, scales, zero_point = _quantize_blockwise_int4(lm_head_weight.astype(np.float32), block_size)
+
+    input_ids = helper.make_tensor_value_info("input_ids", TensorProto.INT64, [batch, seq])
+    hidden_states = helper.make_tensor_value_info("hidden_states", TensorProto.FLOAT16, [batch, seq, hidden])
+    logits = helper.make_tensor_value_info("logits", TensorProto.FLOAT16, [batch, seq, vocab])
+    embeds = helper.make_tensor_value_info("embeds", TensorProto.FLOAT16, [batch, seq, hidden])
+
+    initializers = [
+        numpy_helper.from_array(qweight, name="embed.weight_Q4"),
+        numpy_helper.from_array(scales, name="embed.weight_scales"),
+        numpy_helper.from_array(zero_point, name="embed.weight_zero_point"),
+        numpy_helper.from_array(lm_head_weight, name="embed.weight"),
+    ]
+    nodes = [
+        helper.make_node(
+            "GatherBlockQuantized",
+            ["embed.weight_Q4", "input_ids", "embed.weight_scales", "embed.weight_zero_point"],
+            ["embeds"],
+            name="embed/GatherBlockQuantized",
+            domain=MSFT_DOMAIN,
+            gather_axis=0,
+            quantize_axis=1,
+            block_size=block_size,
+            bits=4,
+        ),
+        helper.make_node("Transpose", ["embed.weight"], ["lm_head.weight_t"], name="lm_head/Transpose", perm=[1, 0]),
+        helper.make_node("MatMul", ["hidden_states", "lm_head.weight_t"], ["logits"], name="lm_head/MatMul"),
+    ]
+    graph = helper.make_graph(
+        nodes=nodes,
+        name="TiedReuseGraph",
+        inputs=[input_ids, hidden_states],
+        outputs=[logits, embeds],
+        initializer=initializers,
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid(MSFT_DOMAIN, 1)])
+    model.ir_version = 10
+    model_path = tmp_path / "model.onnx"
+    onnx.save(model, model_path)
+    return ONNXModelHandler(model_path=str(model_path))
+
+
+def test_tie_word_embeddings_reuse_shares_quantized_embedding(tmp_path):
+    # setup: LM head weight equals the (quantized) embedding table -> tied
+    np.random.seed(0)
+    weight = np.random.randn(6, 8).astype(np.float16)
+    input_model = _build_quantized_embedding_tied_model(tmp_path, weight)
+
+    p = create_pass_from_dict(GraphSurgeries, {"surgeries": [{"surgeon": "TieWordEmbeddings"}]}, disable_search=True)
+
+    # execute
+    output_model = p.run(input_model, str(tmp_path / "output"))
+
+    # assert: the float MatMul LM head is now a MatMulNBits sharing the embedding's
+    # quantized tensors, and the redundant Transpose + float weight are pruned.
+    dag = OnnxDAG.from_model_path(output_model.model_path)
+    op_types = [dag.get_node_op_type(n) for n in dag.get_node_names()]
+    assert "MatMul" not in op_types
+    assert "Transpose" not in op_types
+    assert op_types.count("MatMulNBits") == 1
+
+    logits_producer = dag.get_producer("logits")
+    assert dag.get_node_op_type(logits_producer) == "MatMulNBits"
+    # scales / zero_point are shared with the embedding node
+    lm_head_inputs = dag.get_node_inputs(logits_producer)
+    assert "embed.weight_scales" in lm_head_inputs
+    assert "embed.weight_zero_point" in lm_head_inputs
+    # the float embedding weight is no longer serialized
+    assert "embed.weight" not in dag.ios
+
+
+def test_tie_word_embeddings_reuse_skips_untied_lm_head(tmp_path):
+    # setup: LM head weight is unrelated to the embedding table -> must NOT tie
+    np.random.seed(0)
+    embed_weight = np.random.randn(6, 8).astype(np.float16)
+    unrelated = (embed_weight + 5.0).astype(np.float16)
+    # quantize/serialize using the embedding weight, but feed an unrelated float weight
+    input_model = _build_quantized_embedding_tied_model(tmp_path, embed_weight)
+    # rewrite the float lm_head weight to an unrelated tensor
+    model = onnx.load(input_model.model_path)
+    for init in model.graph.initializer:
+        if init.name == "embed.weight":
+            init.CopyFrom(numpy_helper.from_array(unrelated, name="embed.weight"))
+    onnx.save(model, input_model.model_path)
+
+    p = create_pass_from_dict(GraphSurgeries, {"surgeries": [{"surgeon": "TieWordEmbeddings"}]}, disable_search=True)
+
+    # execute
+    output_model = p.run(input_model, str(tmp_path / "output"))
+
+    # assert: gate rejects the untied weight, LM head stays a float MatMul
+    dag = OnnxDAG.from_model_path(output_model.model_path)
+    assert dag.get_node_op_type(dag.get_producer("logits")) == "MatMul"
+
+
 def test_remove_gidx_from_matmulnbits(tmp_path):
     # setup
     a_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3])
