SHARP-onnx / quantize.py

Cool2009

Upload folder using huggingface_hub

7871e64 verified 1 day ago

6.59 kB

"""Generate the reduced-precision variants of the SHARP ONNX model.
Each variant is written to its own folder next to `fp32/`, as `model.onnx` plus
`model.onnx.data` (the external weight file).
    python quantize.py fp16
    python quantize.py int8
    python quantize.py int4
    python quantize.py fp8
Requires: onnx, onnxruntime, onnxconverter-common (fp16), ml_dtypes (fp8).
"""

import argparse
import os

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper

SRC_MODEL = os.path.join("fp32", "model.onnx")
DATA_FILE = "model.onnx.data"


def save_model(model: onnx.ModelProto, folder: str) -> str:
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, "model.onnx")
    # onnx.save appends to an existing external-data file instead of replacing it,
    # so any leftovers from an earlier run have to go first
    for stale in (path, os.path.join(folder, DATA_FILE)):
        if os.path.exists(stale):
            os.remove(stale)
    onnx.save(
        model,
        path,
        save_as_external_data=True,
        all_tensors_to_one_file=True,
        location=DATA_FILE,
        size_threshold=1024,
        convert_attribute=False,
    )
    return path


def _cast_resize_fp16_inputs(model: onnx.ModelProto) -> int:
    """Insert Cast nodes so Resize gets float32 roi/scales.
    Operator block lists do not apply to `Constant` nodes folded into the graph, so the
    fp16 scales have to be cast back explicitly.
    """
    initializers = {t.name: t for t in model.graph.initializer}
    producers = {out: node for node in model.graph.node for out in node.output}
    patched = 0
    nodes = []
    for node in model.graph.node:
        if node.op_type != "Resize":
            nodes.append(node)
            continue
        for index in range(1, len(node.input)):
            name = node.input[index]
            if not name:
                continue
            producer = producers.get(name)
            is_fp16 = initializers.get(name) is not None and initializers[name].data_type == TensorProto.FLOAT16
            if producer is not None and producer.op_type == "Constant":
                tensor = [a.t for a in producer.attribute if a.name == "value"][0]
                is_fp16 = tensor.data_type == TensorProto.FLOAT16
            if not is_fp16:
                continue
            cast_out = f"{name}_float32"
            nodes.append(helper.make_node("Cast", [name], [cast_out], name=f"{cast_out}_cast", to=TensorProto.FLOAT))
            node.input[index] = cast_out
            patched += 1
        nodes.append(node)
    if patched:
        model.graph.ClearField("node")
        model.graph.node.extend(nodes)
    return patched


def convert_fp16() -> str:
    # onnxruntime's converter inserts the casts needed to keep every node's operand
    # dtypes consistent, which onnxconverter-common does not do reliably.
    from onnxruntime.transformers.float16 import convert_float_to_float16

    model = onnx.load(SRC_MODEL)  # resolves the external .data file
    # use_symbolic_shape_infer=False: the in-memory proto is >2 GB and cannot be serialized
    model = convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=True)
    print(f"fp16: cast {_cast_resize_fp16_inputs(model)} Resize inputs back to float32")
    return save_model(model, "fp16")


def convert_int8() -> str:
    from onnxruntime.quantization import QuantType, quantize_dynamic

    out = os.path.join("int8", "model.onnx")
    os.makedirs("int8", exist_ok=True)
    quantize_dynamic(
        SRC_MODEL,
        out,
        weight_type=QuantType.QInt8,
        per_channel=True,
        reduce_range=False,
        use_external_data_format=True,
        extra_options={"MatMulConstBOnly": False},
    )
    return out


def convert_int4() -> str:
    from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer

    model = onnx.load(SRC_MODEL)
    quantizer = MatMulNBitsQuantizer(model, block_size=32, is_symmetric=True, bits=4)
    quantizer.process()
    # newer onnxruntime wraps the ModelProto in onnx_ir.Model
    quantized = quantizer.model.model if hasattr(quantizer.model, "model") else quantizer.model
    return save_model(quantized, "int4")


def convert_fp8() -> str:
    """Store weights as float8_e4m3fn with a per-output-channel scale.
    Each weight becomes `<name>`_fp8` plus a float32 scale, and a DequantizeLinear node
    restores `<name>`. float8 needs opset 19 and ONNX Runtime >= 1.25 for the CPU kernel.
    """
    import ml_dtypes

    model = onnx.load(SRC_MODEL)
    for opset in model.opset_import:
        if opset.domain in ("", "ai.onnx"):
            opset.version = max(opset.version, 19)

    dequantize = []
    new_initializers = []
    converted = 0
    for init in model.graph.initializer:
        array = numpy_helper.to_array(init)
        if array.dtype != np.float32 or array.ndim < 2 or array.size < 1024:
            continue
        name = init.name
        amax = np.abs(array).reshape(array.shape[0], -1).max(axis=1)
        scale = np.maximum(amax, 1e-9).astype(np.float32) / 448.0  # 448 = E4M3 max
        shape = (array.shape[0],) + (1,) * (array.ndim - 1)
        quantized = (array / scale.reshape(shape)).astype(ml_dtypes.float8_e4m3fn)
        # the fp8 tensor takes a new name; DequantizeLinear rebuilds the original one
        init.CopyFrom(numpy_helper.from_array(quantized, name + "_fp8"))
        new_initializers.append(numpy_helper.from_array(scale, name + "_scale"))
        dequantize.append(
            helper.make_node("DequantizeLinear", [name + "_fp8", name + "_scale"], [name],
                             name=f"{name}_dequantize", axis=0)
        )
        converted += 1

    # the DequantizeLinear nodes read only initializers, so they can run first
    nodes = list(dequantize) + list(model.graph.node)
    model.graph.ClearField("node")
    model.graph.node.extend(nodes)
    model.graph.initializer.extend(new_initializers)
    print(f"fp8: converted {converted} weight tensors")
    return save_model(model, "fp8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("variant", choices=["fp16", "int8", "int4", "fp8", "all"])
    args = parser.parse_args()

    steps = (
        [convert_fp16, convert_int8, convert_int4, convert_fp8]
        if args.variant == "all"
        else [globals()[f"convert_{args.variant}"]]
    )
    for step in steps:
        print(f"-> {step.__name__}")
        print("   ", step())


if __name__ == "__main__":
    main()

	"""Generate the reduced-precision variants of the SHARP ONNX model.

	Each variant is written to its own folder next to `fp32/`, as `model.onnx` plus
	`model.onnx.data` (the external weight file).

	python quantize.py fp16
	python quantize.py int8
	python quantize.py int4
	python quantize.py fp8

	Requires: onnx, onnxruntime, onnxconverter-common (fp16), ml_dtypes (fp8).
	"""

	import argparse
	import os

	import numpy as np
	import onnx
	from onnx import TensorProto, helper, numpy_helper

	SRC_MODEL = os.path.join("fp32", "model.onnx")
	DATA_FILE = "model.onnx.data"


	def save_model(model: onnx.ModelProto, folder: str) -> str:
	os.makedirs(folder, exist_ok=True)
	path = os.path.join(folder, "model.onnx")
	# onnx.save appends to an existing external-data file instead of replacing it,
	# so any leftovers from an earlier run have to go first
	for stale in (path, os.path.join(folder, DATA_FILE)):
	if os.path.exists(stale):
	os.remove(stale)
	onnx.save(
	model,
	path,
	save_as_external_data=True,
	all_tensors_to_one_file=True,
	location=DATA_FILE,
	size_threshold=1024,
	convert_attribute=False,
	)
	return path


	def _cast_resize_fp16_inputs(model: onnx.ModelProto) -> int:
	"""Insert Cast nodes so Resize gets float32 roi/scales.

	Operator block lists do not apply to `Constant` nodes folded into the graph, so the
	fp16 scales have to be cast back explicitly.
	"""
	initializers = {t.name: t for t in model.graph.initializer}
	producers = {out: node for node in model.graph.node for out in node.output}
	patched = 0
	nodes = []
	for node in model.graph.node:
	if node.op_type != "Resize":
	nodes.append(node)
	continue
	for index in range(1, len(node.input)):
	name = node.input[index]
	if not name:
	continue
	producer = producers.get(name)
	is_fp16 = initializers.get(name) is not None and initializers[name].data_type == TensorProto.FLOAT16
	if producer is not None and producer.op_type == "Constant":
	tensor = [a.t for a in producer.attribute if a.name == "value"][0]
	is_fp16 = tensor.data_type == TensorProto.FLOAT16
	if not is_fp16:
	continue
	cast_out = f"{name}_float32"
	nodes.append(helper.make_node("Cast", [name], [cast_out], name=f"{cast_out}_cast", to=TensorProto.FLOAT))
	node.input[index] = cast_out
	patched += 1
	nodes.append(node)
	if patched:
	model.graph.ClearField("node")
	model.graph.node.extend(nodes)
	return patched


	def convert_fp16() -> str:
	# onnxruntime's converter inserts the casts needed to keep every node's operand
	# dtypes consistent, which onnxconverter-common does not do reliably.
	from onnxruntime.transformers.float16 import convert_float_to_float16

	model = onnx.load(SRC_MODEL) # resolves the external .data file
	# use_symbolic_shape_infer=False: the in-memory proto is >2 GB and cannot be serialized
	model = convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=True)
	print(f"fp16: cast {_cast_resize_fp16_inputs(model)} Resize inputs back to float32")
	return save_model(model, "fp16")


	def convert_int8() -> str:
	from onnxruntime.quantization import QuantType, quantize_dynamic

	out = os.path.join("int8", "model.onnx")
	os.makedirs("int8", exist_ok=True)
	quantize_dynamic(
	SRC_MODEL,
	out,
	weight_type=QuantType.QInt8,
	per_channel=True,
	reduce_range=False,
	use_external_data_format=True,
	extra_options={"MatMulConstBOnly": False},
	)
	return out


	def convert_int4() -> str:
	from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer

	model = onnx.load(SRC_MODEL)
	quantizer = MatMulNBitsQuantizer(model, block_size=32, is_symmetric=True, bits=4)
	quantizer.process()
	# newer onnxruntime wraps the ModelProto in onnx_ir.Model
	quantized = quantizer.model.model if hasattr(quantizer.model, "model") else quantizer.model
	return save_model(quantized, "int4")


	def convert_fp8() -> str:
	"""Store weights as float8_e4m3fn with a per-output-channel scale.

	Each weight becomes `<name>`_fp8` plus a float32 scale, and a DequantizeLinear node
	restores `<name>`. float8 needs opset 19 and ONNX Runtime >= 1.25 for the CPU kernel.
	"""
	import ml_dtypes

	model = onnx.load(SRC_MODEL)
	for opset in model.opset_import:
	if opset.domain in ("", "ai.onnx"):
	opset.version = max(opset.version, 19)

	dequantize = []
	new_initializers = []
	converted = 0
	for init in model.graph.initializer:
	array = numpy_helper.to_array(init)
	if array.dtype != np.float32 or array.ndim < 2 or array.size < 1024:
	continue
	name = init.name
	amax = np.abs(array).reshape(array.shape[0], -1).max(axis=1)
	scale = np.maximum(amax, 1e-9).astype(np.float32) / 448.0 # 448 = E4M3 max
	shape = (array.shape[0],) + (1,) * (array.ndim - 1)
	quantized = (array / scale.reshape(shape)).astype(ml_dtypes.float8_e4m3fn)
	# the fp8 tensor takes a new name; DequantizeLinear rebuilds the original one
	init.CopyFrom(numpy_helper.from_array(quantized, name + "_fp8"))
	new_initializers.append(numpy_helper.from_array(scale, name + "_scale"))
	dequantize.append(
	helper.make_node("DequantizeLinear", [name + "_fp8", name + "_scale"], [name],
	name=f"{name}_dequantize", axis=0)
	)
	converted += 1

	# the DequantizeLinear nodes read only initializers, so they can run first
	nodes = list(dequantize) + list(model.graph.node)
	model.graph.ClearField("node")
	model.graph.node.extend(nodes)
	model.graph.initializer.extend(new_initializers)
	print(f"fp8: converted {converted} weight tensors")
	return save_model(model, "fp8")


	def main() -> None:
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("variant", choices=["fp16", "int8", "int4", "fp8", "all"])
	args = parser.parse_args()

	steps = (
	[convert_fp16, convert_int8, convert_int4, convert_fp8]
	if args.variant == "all"
	else [globals()[f"convert_{args.variant}"]]
	)
	for step in steps:
	print(f"-> {step.__name__}")
	print(" ", step())


	if __name__ == "__main__":
	main()