Reuse quantized embedding table for tied LM head in TieWordEmbeddings

[justinchuby]

Describe your changes

Add a reuse mode to the TieWordEmbeddings graph surgery.

Previously TieWordEmbeddings handled two cases: both weights unquantized
(Gather + MatMul) or both quantized (GatherBlockQuantized + MatMulNBits).
There was no path for when only the embedding is quantized while the LM head
is still a float MatMul whose weight is the tied embedding (reached through a
Transpose). This happens naturally when the embedding Gather is quantized with
OnnxBlockWiseRtnQuantization while the transformer body is left for a separate
pass such as OnnxKQuantQuantization. In that state the tied word-embedding matrix
is stored twice — once as INT4 (embedding) and once as float16 (LM head) —
which is larger than an all-float16 model.

This PR adds:

• handle_reuse — rebuilds the float LM head as a MatMulNBits that shares the
  embedding's INT4 qweight / scales / zero_point (the byte-identical table,
  reshaped from the 2D GatherBlockQuantized layout to the 3D MatMulNBits
  layout), then prunes the now-dead Transpose and the float embedding weight.
• reuse_weights_match — a correctness gate that only ties when the float LM
  head weight actually equals the dequantized embedding table (comparing a slice),
  so an untied projection is never incorrectly tied.

Pipeline it enables (smallest size at the highest-quality body quantization):

MobiusBuilder(fp16)
  -> OnnxKQuantQuantization            # body -> Q4_K_M
  -> OnnxBlockWiseRtnQuantization      # embedding Gather -> GatherBlockQuantized
  -> GraphSurgeries[TieWordEmbeddings] # LM head reuses the embedding INT4 table

Each pass touches only its intended nodes (body MatMuls have initializer weights;
the embedding Gather has an initializer weight; the tied LM head's weight is
behind a Transpose, so it is skipped by both quantizers and handled here).

Measured on a 1.8B tied-embedding translation model (CUDA): on-disk size drops from
~1.39 GB (K-Quant body, float16 embedding/LM head) to ~1.03 GB, with equal or
better output fidelity vs float16 compared to the two-table INT4 variant.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.
  TieWordEmbeddings can now tie a float LM head onto an already-quantized
  (GatherBlockQuantized) embedding, storing the shared word-embedding matrix
  once as INT4 instead of INT4 + float16.

(Optional) Issue link

[Copilot]

Pull request overview

This PR extends the TieWordEmbeddings ONNX graph surgery to support a new reuse mode where the embedding is already quantized (GatherBlockQuantized) but the tied LM head is still a float MatMul. In this case, the surgery rebuilds the LM head as MatMulNBits that reuses the embedding’s quantized tensors to avoid storing the same embedding table twice.

Changes:

• Add a reuse path in TieWordEmbeddings to convert a float LM-head MatMul into MatMulNBits sharing the embedding’s quantized qweight/scales/zero_point.
• Add a reuse_weights_match gate to verify the float LM-head weight matches the dequantized embedding table before tying.
• Add unit tests covering both the successful reuse case and the “skip when not actually tied” case.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
olive/passes/onnx/graph_surgeries.py	Adds the reuse-mode surgery + correctness gate and performs pruning/rewiring of the LM head to reuse embedding quantized tensors.
test/passes/onnx/test_graph_surgeries.py	Adds unit tests that validate reuse-mode tying occurs only when the float LM head truly matches the embedding table.