Postprocessing to share lm_head weights to embedding

[jiafatom]

This PR is based on #1437, but we don't need convert MatMul here.

[tianleiwu]

Can we use DequantizeLinear op?
https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html

[jiafatom]

Could you please elaborate more, how to construct?

[tianleiwu]

Use helper.make_node to create a DequantizeLinear node, and feed the quantized lm weight and same scale and bias used in last MatMulNBits node into DequantizeLinear. Then you can get the dequantized weights. Then you can Gather based on input_ids.

[kunal-vaishnavi]

I agree that we should use DequantizeLinear. It is already constructed in the model builder.

onnxruntime-genai/src/python/py/models/builder.py

Lines 861 to 886 in 8eed730

	def make_dequantize_linear(self, dequantize_name, quantized_op):
	# Input weights are quantized, save quantized MatMul weights for onnx model
	qweight = dequantize_name[1:].replace("/", ".") + ".qweight"
	qweight_npy = quantized_op.qweight.detach().cpu()
	qweight_npy = qweight_npy.reshape(*qweight_npy.shape[:-2], qweight_npy.shape[-2] * qweight_npy.shape[-1])
	self.make_external_tensor(qweight_npy.contiguous(), qweight, True)
	scales = dequantize_name[1:].replace("/", ".") + ".scales"
	scales_npy = quantized_op.scales.detach().cpu().to(self.to_torch_dtype[self.io_dtype])
	scales_npy = scales_npy.reshape(*qweight_npy.shape[:-1], qweight_npy.shape[-1] * 2 // quantized_op.group_size)
	self.make_external_tensor(scales_npy.contiguous(), scales)
	dequantize_inputs = [qweight, scales]
	if hasattr(quantized_op, "qzeros") and quantized_op.qzeros is not None:
	zeros = dequantize_name[1:].replace("/", ".") + ".qzeros"
	zeros_npy = quantized_op.qzeros.detach().cpu()
	zeros_npy = zeros_npy.reshape(*qweight_npy.shape[:-1], qweight_npy.shape[-1] // quantized_op.group_size)
	self.make_external_tensor(zeros_npy.contiguous(), zeros, True)
	dequantize_inputs.append(zeros)
	dequantize_output = f"{dequantize_name}/output_0"
	self.make_node("DequantizeLinear", inputs=dequantize_inputs, outputs=[dequantize_output], name=dequantize_name, block_size=quantized_op.group_size, axis=-1)
	self.make_value_info(dequantize_output, self.io_dtype, shape=[*scales_npy.shape[:-1], scales_npy.shape[-1] * quantized_op.group_size])
	return dequantize_output

It will also be easier to construct the temporary subgraph for GatherBlockQuantized in the model builder directly.