1132 lines
54 KiB
Python
1132 lines
54 KiB
Python
|
# mypy: ignore-errors
|
||
|
|
|
||
|
|
from typing import Any, Dict, List, Optional, Set, Tuple, Union, Type, Callable
|
||
|
|
from torch.ao.quantization.quant_type import QuantType
|
||
|
|
import torch
|
||
|
|
import copy
|
||
|
|
import warnings
|
||
|
|
from torch.fx import (
|
||
|
|
GraphModule,
|
||
|
|
)
|
||
|
|
from torch.fx.graph import (
|
||
|
|
Graph,
|
||
|
|
Node,
|
||
|
|
Argument,
|
||
|
|
)
|
||
|
|
from ..utils import (
|
||
|
|
activation_is_statically_quantized,
|
||
|
|
weight_is_quantized,
|
||
|
|
get_qparam_dict,
|
||
|
|
_parent_name,
|
||
|
|
get_swapped_custom_module_class,
|
||
|
|
)
|
||
|
|
from ..qconfig import (
|
||
|
|
QConfigAny,
|
||
|
|
qconfig_equals
|
||
|
|
)
|
||
|
|
from ..qconfig_mapping import QConfigMapping
|
||
|
|
from .qconfig_mapping_utils import (
|
||
|
|
_generate_node_name_to_qconfig,
|
||
|
|
_compare_prepare_convert_qconfig_mappings,
|
||
|
|
_update_qconfig_for_fusion,
|
||
|
|
_is_qconfig_supported_by_dtype_configs,
|
||
|
|
_update_qconfig_for_qat,
|
||
|
|
)
|
||
|
|
from torch.ao.quantization.backend_config.utils import (
|
||
|
|
get_root_module_to_quantized_reference_module,
|
||
|
|
get_pattern_to_dtype_configs,
|
||
|
|
get_fused_module_classes,
|
||
|
|
get_qat_module_classes,
|
||
|
|
)
|
||
|
|
from torch.ao.quantization.backend_config import (
|
||
|
|
BackendConfig,
|
||
|
|
get_native_backend_config,
|
||
|
|
)
|
||
|
|
from torch.ao.quantization.observer import _is_activation_post_process
|
||
|
|
from .graph_module import (
|
||
|
|
_is_observed_module,
|
||
|
|
_is_observed_standalone_module,
|
||
|
|
)
|
||
|
|
from ._equalize import update_obs_for_equalization, convert_eq_obs
|
||
|
|
from torch.nn.utils.parametrize import type_before_parametrizations
|
||
|
|
from .utils import (
|
||
|
|
_get_module,
|
||
|
|
_is_custom_module_lstm,
|
||
|
|
_is_custom_module_mha,
|
||
|
|
assert_and_get_unique_device,
|
||
|
|
get_custom_module_class_keys,
|
||
|
|
create_getattr_from_value,
|
||
|
|
collect_producer_nodes,
|
||
|
|
graph_module_from_producer_nodes,
|
||
|
|
node_arg_is_weight,
|
||
|
|
)
|
||
|
|
from torch.ao.quantization.utils import (
|
||
|
|
is_per_channel,
|
||
|
|
to_underlying_dtype,
|
||
|
|
)
|
||
|
|
from torch.ao.quantization.quantize import (
|
||
|
|
_remove_qconfig,
|
||
|
|
)
|
||
|
|
from torch.ao.quantization.stubs import DeQuantStub
|
||
|
|
from .custom_config import (
|
||
|
|
ConvertCustomConfig,
|
||
|
|
PrepareCustomConfig,
|
||
|
|
)
|
||
|
|
from .lower_to_fbgemm import lower_to_fbgemm
|
||
|
|
# importing the lib so that the quantized_decomposed ops are registered
|
||
|
|
from ._decomposed import quantized_decomposed_lib # noqa: F401
|
||
|
|
import operator
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"convert",
|
||
|
|
"convert_custom_module",
|
||
|
|
"convert_standalone_module",
|
||
|
|
"convert_weighted_module",
|
||
|
|
]
|
||
|
|
|
||
|
|
_QSCHEME_TO_CHOOSE_QPARAMS_OP = {
|
||
|
|
torch.per_tensor_affine: torch.ops.quantized_decomposed.choose_qparams.tensor,
|
||
|
|
torch.per_tensor_symmetric: torch.ops.quantized_decomposed.choose_qparams_symmetric.tensor,
|
||
|
|
}
|
||
|
|
|
||
|
|
def _replace_observer_with_quantize_dequantize_node_decomposed(
|
||
|
|
model: torch.fx.GraphModule,
|
||
|
|
node: Node,
|
||
|
|
modules: Dict[str, torch.nn.Module],
|
||
|
|
node_name_to_scope: Dict[str, Tuple[str, type]],
|
||
|
|
node_name_to_qconfig: Dict[str, QConfigAny]) -> None:
|
||
|
|
""" Replace activation_post_process module call node with quantize and
|
||
|
|
dequantize node working with decomposed Tensor
|
||
|
|
|
||
|
|
Before:
|
||
|
|
... -> observer_0(x) -> ...
|
||
|
|
After:
|
||
|
|
... -> torch.ops.quantized_decomposed.quantize_per_tensor(x, ...) ->
|
||
|
|
torch.ops.quantized_decomposed.dequantize_per_tensor() -> ...
|
||
|
|
|
||
|
|
or quantize_per_channel and dequantize_per_channel
|
||
|
|
"""
|
||
|
|
graph = model.graph
|
||
|
|
assert modules is not None
|
||
|
|
assert isinstance(node.target, str)
|
||
|
|
module_path, prefix = _get_module_path_and_prefix(node, node_name_to_scope, node_name_to_qconfig)
|
||
|
|
activation_post_process = modules[node.target]
|
||
|
|
if hasattr(activation_post_process, "convert"):
|
||
|
|
activation_post_process.convert(model, node)
|
||
|
|
return
|
||
|
|
# skip replacing observers to quant/dequant nodes if the qconfigs of all
|
||
|
|
# consumers and producers of this observer are None
|
||
|
|
skip_replacement = all(_has_none_qconfig(n, node_name_to_qconfig) for n in
|
||
|
|
list(node.args) + list(node.users.keys()))
|
||
|
|
if skip_replacement or not _is_conversion_supported(activation_post_process):
|
||
|
|
# didn't find corresponding quantize op and info for the activation_post_process
|
||
|
|
# so we just remove the observer
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
node.replace_all_uses_with(node.args[0])
|
||
|
|
graph.erase_node(node)
|
||
|
|
return
|
||
|
|
|
||
|
|
# otherwise, we can convert the activation_post_process module call to quantize/dequantize node
|
||
|
|
|
||
|
|
# 1. extract the information from activation_post_process module for generating
|
||
|
|
# the quantize and dequantize operator
|
||
|
|
dtype = activation_post_process.dtype # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
is_dynamic = False
|
||
|
|
if hasattr(activation_post_process, "is_dynamic"):
|
||
|
|
is_dynamic = activation_post_process.is_dynamic # type: ignore[assignment]
|
||
|
|
|
||
|
|
if dtype in [torch.quint8, torch.qint8, torch.qint32, torch.uint8, torch.int8, torch.int16, torch.int32] and \
|
||
|
|
(not is_dynamic):
|
||
|
|
# TODO: probably should cleanup this condition check, it's hard
|
||
|
|
# to reason about this if and the following elif
|
||
|
|
|
||
|
|
# uint8/int8/int32 static quantization branch
|
||
|
|
|
||
|
|
# 1. extract information for inserting q/dq node from activation_post_process
|
||
|
|
node_type = "call_function"
|
||
|
|
quantize_op : Optional[Callable] = None
|
||
|
|
scale, zero_point = activation_post_process.calculate_qparams() # type: ignore[attr-defined, operator]
|
||
|
|
if is_per_channel(activation_post_process.qscheme): # type: ignore[attr-defined]
|
||
|
|
ch_axis = int(activation_post_process.ch_axis) # type: ignore[attr-defined, arg-type]
|
||
|
|
quantize_op = torch.ops.quantized_decomposed.quantize_per_channel.default
|
||
|
|
dequantize_op = torch.ops.quantized_decomposed.dequantize_per_channel.default
|
||
|
|
quant_min = activation_post_process.quant_min
|
||
|
|
quant_max = activation_post_process.quant_max
|
||
|
|
dtype_ = to_underlying_dtype(dtype)
|
||
|
|
qparams = {
|
||
|
|
"_scale_": scale,
|
||
|
|
"_zero_point_": zero_point,
|
||
|
|
"_axis_": ch_axis,
|
||
|
|
"_quant_min_": quant_min,
|
||
|
|
"_quant_max_": quant_max,
|
||
|
|
"_dtype_": dtype_
|
||
|
|
}
|
||
|
|
else:
|
||
|
|
quantize_op = torch.ops.quantized_decomposed.quantize_per_tensor.default
|
||
|
|
dequantize_op = torch.ops.quantized_decomposed.dequantize_per_tensor.default
|
||
|
|
scale = float(scale)
|
||
|
|
zero_point = int(zero_point)
|
||
|
|
quant_min = activation_post_process.quant_min # type: ignore[attr-defined]
|
||
|
|
quant_max = activation_post_process.quant_max # type: ignore[attr-defined]
|
||
|
|
dtype_ = to_underlying_dtype(dtype)
|
||
|
|
qparams = {
|
||
|
|
"_scale_": scale,
|
||
|
|
"_zero_point_": zero_point,
|
||
|
|
"_quant_min_": quant_min,
|
||
|
|
"_quant_max_": quant_max,
|
||
|
|
"_dtype_": dtype_
|
||
|
|
}
|
||
|
|
|
||
|
|
# 2. replace activation_post_process node with quantize and dequantize
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
input_node = node.args[0]
|
||
|
|
quantize_op_inputs = [input_node]
|
||
|
|
for key, value_or_node in qparams.items():
|
||
|
|
# TODO: we can add the information of whether a value needs to
|
||
|
|
# be registered as an attribute in qparams dict itself
|
||
|
|
if key in ['_scale_', '_zero_point_'] and (not isinstance(value_or_node, (float, int))):
|
||
|
|
# For scale and zero_point values we register them as buffers in the root module.
|
||
|
|
# However, note that when the values are not tensors, as in the case of
|
||
|
|
# per_tensor quantization, they will be treated as literals.
|
||
|
|
# However, registering them as a node seems to cause issue with dynamo
|
||
|
|
# tracing where it may consider tensor overload as opposed to default.
|
||
|
|
# With extra check of scale and zero_point being scalar, it makes
|
||
|
|
# sure that the default overload can be used.
|
||
|
|
# TODO: maybe need more complex attr name here
|
||
|
|
qparam_node = create_getattr_from_value(
|
||
|
|
model, graph, module_path + prefix + key, value_or_node)
|
||
|
|
quantize_op_inputs.append(qparam_node)
|
||
|
|
else:
|
||
|
|
# for qparams that are not scale/zero_point (like axis, dtype) we store them as literals in the graph.
|
||
|
|
quantize_op_inputs.append(value_or_node)
|
||
|
|
|
||
|
|
quantized_node = graph.create_node(node_type, quantize_op, tuple(quantize_op_inputs), {})
|
||
|
|
# use the same qparams from quantize op
|
||
|
|
dq_inputs = [quantized_node] + quantize_op_inputs[1:]
|
||
|
|
dequantized_node = graph.call_function(
|
||
|
|
dequantize_op,
|
||
|
|
tuple(dq_inputs),
|
||
|
|
{}
|
||
|
|
)
|
||
|
|
|
||
|
|
def remap_fn(x):
|
||
|
|
return dequantized_node if x is node else x
|
||
|
|
|
||
|
|
# remap numeric_debug_handle
|
||
|
|
for user_node in node.users:
|
||
|
|
if "numeric_debug_handle" in user_node.meta:
|
||
|
|
numeric_debug_handle = user_node.meta["numeric_debug_handle"]
|
||
|
|
user_node.meta["numeric_debug_handle"] = {remap_fn(k): v for k, v in numeric_debug_handle.items()}
|
||
|
|
node.replace_all_uses_with(dequantized_node)
|
||
|
|
graph.erase_node(node)
|
||
|
|
elif is_dynamic:
|
||
|
|
|
||
|
|
# uint8/int8/fp16 dynamic quantization
|
||
|
|
|
||
|
|
# 1. extract information for inserting q/dq node from activation_post_process
|
||
|
|
node_type = "call_function"
|
||
|
|
quantize_op = torch.ops.quantized_decomposed.quantize_per_tensor.tensor
|
||
|
|
# we only use choose_qparams for is_decomposed now,
|
||
|
|
# but we should probably align the non-decomposed path with this as well,
|
||
|
|
# and that can be done after we remove reduce_range flag
|
||
|
|
# 1. extract qparams from activation_post_process module
|
||
|
|
dtype_ = to_underlying_dtype(dtype)
|
||
|
|
assert dtype_ in [torch.uint8, torch.int8], \
|
||
|
|
"only uint8 and int8 are supported in reference flow for " \
|
||
|
|
"dynamic quantization right now"
|
||
|
|
quant_min = activation_post_process.quant_min # type: ignore[attr-defined]
|
||
|
|
quant_max = activation_post_process.quant_max # type: ignore[attr-defined]
|
||
|
|
qscheme = getattr(activation_post_process, "qscheme", torch.per_tensor_affine) # type: ignore[attr-defined]
|
||
|
|
eps = getattr(activation_post_process, "eps", torch.finfo(torch.float32).eps) # type: ignore[attr-defined]
|
||
|
|
# note: scale and zero_point are missing for quantize_per_tensor op
|
||
|
|
# we'll need to get this from choose_qparams op, which we'll add after
|
||
|
|
# this step
|
||
|
|
qparams = {
|
||
|
|
"_quant_min_": quant_min,
|
||
|
|
"_quant_max_": quant_max,
|
||
|
|
"_eps_": eps,
|
||
|
|
"_dtype_": dtype_
|
||
|
|
}
|
||
|
|
|
||
|
|
choose_qparams_op = _QSCHEME_TO_CHOOSE_QPARAMS_OP[qscheme]
|
||
|
|
# 2. insert choose_qparams op and update the qparams list
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
input_node = node.args[0]
|
||
|
|
choose_qparams_op_inputs = [node.args[0]]
|
||
|
|
for key, value in qparams.items():
|
||
|
|
# we have quant_min, quant_max and dtype, all should be stored
|
||
|
|
# as literals
|
||
|
|
choose_qparams_op_inputs.append(value)
|
||
|
|
choose_qparams_node = graph.create_node(
|
||
|
|
"call_function",
|
||
|
|
choose_qparams_op,
|
||
|
|
tuple(choose_qparams_op_inputs),
|
||
|
|
{}
|
||
|
|
)
|
||
|
|
# choose_qparms returns (scale, zero_point)
|
||
|
|
scale_node = graph.create_node(
|
||
|
|
"call_function",
|
||
|
|
operator.getitem,
|
||
|
|
(choose_qparams_node, 0),
|
||
|
|
{}
|
||
|
|
)
|
||
|
|
zero_point_node = graph.create_node(
|
||
|
|
"call_function",
|
||
|
|
operator.getitem,
|
||
|
|
(choose_qparams_node, 1),
|
||
|
|
{}
|
||
|
|
)
|
||
|
|
quant_min = qparams["_quant_min_"]
|
||
|
|
quant_max = qparams["_quant_max_"]
|
||
|
|
dtype = qparams["_dtype_"]
|
||
|
|
qparams = {
|
||
|
|
"_scale_": scale_node,
|
||
|
|
"_zero_point_": zero_point_node,
|
||
|
|
"_quant_min_": quant_min,
|
||
|
|
"_quant_max_": quant_max,
|
||
|
|
"_dtype_": dtype
|
||
|
|
}
|
||
|
|
|
||
|
|
# 3. replace activation_post_process node to quantize and dequantize node
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
input_node = node.args[0]
|
||
|
|
quantize_op_inputs = [input_node]
|
||
|
|
for key, value_or_node in qparams.items():
|
||
|
|
# TODO: we can add the information of whether a value needs to
|
||
|
|
# be registered as an attribute in qparams dict itself
|
||
|
|
if key in ['_scale_', '_zero_point_']:
|
||
|
|
# in this case we have a node in the graph since it's dynamically
|
||
|
|
# computed from the input, with choose_qparams op
|
||
|
|
qparam_node = value_or_node
|
||
|
|
quantize_op_inputs.append(qparam_node)
|
||
|
|
else:
|
||
|
|
# for qparams that are not scale/zero_point (like axis, dtype) we
|
||
|
|
# store them as literals in the graph.
|
||
|
|
quantize_op_inputs.append(value_or_node)
|
||
|
|
|
||
|
|
quantized_node = graph.create_node(node_type, quantize_op, tuple(quantize_op_inputs), {})
|
||
|
|
# use the same qparams from quantize op
|
||
|
|
dq_inputs = [quantized_node] + quantize_op_inputs[1:]
|
||
|
|
# need to use the tensor variant of this op, since scale and zero_point
|
||
|
|
# from choose_qparam are Tensors, instead of float/int, this is to
|
||
|
|
# prevent these nodes being traced away by downstream systems
|
||
|
|
dequantize_op = torch.ops.quantized_decomposed.dequantize_per_tensor.tensor
|
||
|
|
dequantized_node = graph.call_function(
|
||
|
|
dequantize_op,
|
||
|
|
tuple(dq_inputs),
|
||
|
|
{}
|
||
|
|
)
|
||
|
|
|
||
|
|
def remap_fn(x):
|
||
|
|
return dequantized_node if x is node else x
|
||
|
|
|
||
|
|
# remap numeric_debug_handle
|
||
|
|
for user_node in node.users:
|
||
|
|
if "numeric_debug_handle" in user_node.meta:
|
||
|
|
numeric_debug_handle = user_node.meta["numeric_debug_handle"]
|
||
|
|
user_node.meta["numeric_debug_handle"] = {remap_fn(k): v for k, v in numeric_debug_handle.items()}
|
||
|
|
node.replace_all_uses_with(dequantized_node)
|
||
|
|
graph.erase_node(node)
|
||
|
|
elif dtype == torch.float16:
|
||
|
|
raise NotImplementedError("decomposed to float16 op not implemented yet")
|
||
|
|
|
||
|
|
# should not reach since we have checks in the beginning to make sure the
|
||
|
|
# activation_post_process is supported
|
||
|
|
|
||
|
|
def _replace_observer_with_quantize_dequantize_node(
|
||
|
|
model: torch.fx.GraphModule,
|
||
|
|
node: Node,
|
||
|
|
modules: Dict[str, torch.nn.Module],
|
||
|
|
node_name_to_scope: Dict[str, Tuple[str, type]],
|
||
|
|
node_name_to_qconfig: Dict[str, QConfigAny]) -> None:
|
||
|
|
""" Replace activation_post_process module call node with quantize and
|
||
|
|
dequantize node
|
||
|
|
|
||
|
|
Before:
|
||
|
|
... -> observer_0(x) -> ...
|
||
|
|
After:
|
||
|
|
... -> torch.quantize_per_tensor(x, ...) -> x.dequantize() -> ...
|
||
|
|
"""
|
||
|
|
assert modules is not None
|
||
|
|
assert isinstance(node.target, str)
|
||
|
|
graph = model.graph
|
||
|
|
module_path, prefix = _get_module_path_and_prefix(node, node_name_to_scope, node_name_to_qconfig)
|
||
|
|
activation_post_process = modules[node.target]
|
||
|
|
# skip replacing observers to quant/dequant nodes if the qconfigs of all
|
||
|
|
# consumers and producers of this observer are None
|
||
|
|
skip_replacement = all(_has_none_qconfig(n, node_name_to_qconfig) for n in
|
||
|
|
list(node.args) + list(node.users.keys()))
|
||
|
|
if skip_replacement or not _is_conversion_supported(activation_post_process):
|
||
|
|
# didn't find corresponding quantize op and info for the activation_post_process
|
||
|
|
# so we just remove the observer
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
node.replace_all_uses_with(node.args[0])
|
||
|
|
graph.erase_node(node)
|
||
|
|
return
|
||
|
|
|
||
|
|
# otherwise, we can convert the activation_post_process module call to quantize/dequantize node
|
||
|
|
dtype = activation_post_process.dtype # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
is_dynamic = False
|
||
|
|
if hasattr(activation_post_process, "is_dynamic"):
|
||
|
|
is_dynamic = activation_post_process.is_dynamic # type: ignore[attr-defined, assignment]
|
||
|
|
|
||
|
|
if dtype in [torch.quint8, torch.qint8, torch.qint32] and \
|
||
|
|
(not is_dynamic):
|
||
|
|
# TODO: probably should cleanup this condition check, it's hard
|
||
|
|
# to reason about this if and the following elif
|
||
|
|
|
||
|
|
# uint8/int8/int32 static quantization branch
|
||
|
|
|
||
|
|
# 1. extract the information from activation_post_process module for generating
|
||
|
|
# the quantize and dequantize operator
|
||
|
|
node_type = "call_function"
|
||
|
|
quantize_op : Optional[Callable] = None
|
||
|
|
scale, zero_point = activation_post_process.calculate_qparams() # type: ignore[attr-defined, operator]
|
||
|
|
if is_per_channel(activation_post_process.qscheme): # type: ignore[attr-defined]
|
||
|
|
ch_axis = int(activation_post_process.ch_axis) # type: ignore[attr-defined, arg-type]
|
||
|
|
qparams = {"_scale_": scale, "_zero_point_": zero_point, "_axis_": ch_axis, "_dtype_": dtype}
|
||
|
|
quantize_op = torch.quantize_per_channel
|
||
|
|
else:
|
||
|
|
scale = float(scale)
|
||
|
|
zero_point = int(zero_point)
|
||
|
|
qparams = {"_scale_": scale, "_zero_point_": zero_point, "_dtype_": dtype}
|
||
|
|
quantize_op = torch.quantize_per_tensor
|
||
|
|
|
||
|
|
# 2. replace activation_post_process node with quantize and dequantize
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
input_node = node.args[0]
|
||
|
|
quantize_op_inputs = [input_node]
|
||
|
|
for key, value_or_node in qparams.items():
|
||
|
|
# TODO: we can add the information of whether a value needs to
|
||
|
|
# be registered as an attribute in qparams dict itself
|
||
|
|
if key in ['_scale_', '_zero_point_']:
|
||
|
|
# For scale and zero_point values we register them as buffers in the root module.
|
||
|
|
# TODO: maybe need more complex attr name here
|
||
|
|
qparam_node = create_getattr_from_value(
|
||
|
|
model, graph, module_path + prefix + key, value_or_node)
|
||
|
|
quantize_op_inputs.append(qparam_node)
|
||
|
|
else:
|
||
|
|
# for qparams that are not scale/zero_point (like axis, dtype) we store them as literals in the graph.
|
||
|
|
quantize_op_inputs.append(value_or_node)
|
||
|
|
|
||
|
|
quantized_node = graph.create_node(node_type, quantize_op, tuple(quantize_op_inputs), {})
|
||
|
|
dequantized_node = graph.call_method("dequantize", args=(quantized_node,))
|
||
|
|
node.replace_all_uses_with(dequantized_node)
|
||
|
|
graph.erase_node(node)
|
||
|
|
elif is_dynamic:
|
||
|
|
|
||
|
|
# uint8/int8/fp16 dynamic quantization branch
|
||
|
|
|
||
|
|
node_type = "call_function"
|
||
|
|
quantize_op = torch.quantize_per_tensor_dynamic
|
||
|
|
# TODO: get reduce range from observer
|
||
|
|
# reduce_range = activation_post_process.reduce_range
|
||
|
|
reduce_range = torch.backends.quantized.engine in ("fbgemm", "x86")
|
||
|
|
qparams = {"_dtype_": dtype, "_reduce_range_": reduce_range}
|
||
|
|
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
input_node = node.args[0]
|
||
|
|
quantize_op_inputs = [input_node]
|
||
|
|
for key, value in qparams.items():
|
||
|
|
quantize_op_inputs.append(value)
|
||
|
|
|
||
|
|
quantized_node = graph.create_node(node_type, quantize_op, tuple(quantize_op_inputs), {})
|
||
|
|
dequantized_node = graph.call_method("dequantize", args=(quantized_node,))
|
||
|
|
node.replace_all_uses_with(dequantized_node)
|
||
|
|
graph.erase_node(node)
|
||
|
|
elif dtype == torch.float16:
|
||
|
|
node_type = "call_method"
|
||
|
|
quantize_op = "to" # type: ignore[assignment]
|
||
|
|
qparams = {"_dtype_": dtype}
|
||
|
|
with graph.inserting_before(node):
|
||
|
|
input_node = node.args[0]
|
||
|
|
quantize_op_inputs = [input_node]
|
||
|
|
for key, value in qparams.items():
|
||
|
|
# TODO: we can add the information of whether a value needs to
|
||
|
|
# be registered as an attribute in qparams dict itself
|
||
|
|
quantize_op_inputs.append(value)
|
||
|
|
|
||
|
|
quantized_node = graph.create_node(node_type, quantize_op, tuple(quantize_op_inputs), {})
|
||
|
|
dequantized_node = graph.call_method("dequantize", args=(quantized_node,))
|
||
|
|
node.replace_all_uses_with(dequantized_node)
|
||
|
|
graph.erase_node(node)
|
||
|
|
|
||
|
|
# should not reach since we have checks in the beginning to make sure the
|
||
|
|
# activation_post_process is supported
|
||
|
|
|
||
|
|
# this is a temporary hack for custom module, we may want to implement
|
||
|
|
# this properly after the custom module class design is finalized
|
||
|
|
# TODO: DeQuantStubs are currently inserted only after custom module LSTM, while observers are inserted
|
||
|
|
# after all other custom modules. In the future, we should simply insert QuantStubs before and DeQuantStubs
|
||
|
|
# after custom modules in general, and replace these with "quantize" and "dequantize" nodes respectively.
|
||
|
|
def _replace_observer_or_dequant_stub_with_dequantize_node(node: Node, graph: Graph) -> None:
|
||
|
|
call_custom_module_node = node.args[0]
|
||
|
|
assert isinstance(call_custom_module_node, Node), \
|
||
|
|
f"Expecting the for call custom module node to be a Node, but got {call_custom_module_node}"
|
||
|
|
node.replace_all_uses_with(call_custom_module_node)
|
||
|
|
graph.erase_node(node)
|
||
|
|
_insert_dequantize_node(call_custom_module_node, graph)
|
||
|
|
|
||
|
|
def _is_conversion_supported(activation_post_process: torch.nn.Module) -> bool:
|
||
|
|
dtype = activation_post_process.dtype # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
is_dynamic = False
|
||
|
|
if hasattr(activation_post_process, "is_dynamic"):
|
||
|
|
is_dynamic = activation_post_process.is_dynamic # type: ignore[attr-defined, assignment]
|
||
|
|
|
||
|
|
return (
|
||
|
|
(dtype in [
|
||
|
|
torch.quint8,
|
||
|
|
torch.qint8,
|
||
|
|
torch.qint32,
|
||
|
|
torch.uint8,
|
||
|
|
torch.int8,
|
||
|
|
torch.int16,
|
||
|
|
torch.int32
|
||
|
|
] and (not is_dynamic)) or # type: ignore[return-value]
|
||
|
|
is_dynamic or
|
||
|
|
dtype == torch.float16
|
||
|
|
)
|
||
|
|
|
||
|
|
def _has_none_qconfig(node: Argument, node_name_to_qconfig: Dict[str, QConfigAny]) -> bool:
|
||
|
|
""" Check if a node has a qconfig of None, i.e. user requested to not quantize
|
||
|
|
the node
|
||
|
|
"""
|
||
|
|
return isinstance(node, Node) and node.name in node_name_to_qconfig and node_name_to_qconfig[node.name] is None
|
||
|
|
|
||
|
|
def _run_weight_observers(observed: GraphModule, backend_config: BackendConfig) -> None:
|
||
|
|
""" Extract the subgraph that produces the weight for dynamic quant
|
||
|
|
or weight only quant node and run the subgraph to observe the weight.
|
||
|
|
Note that the observers of dynamic quant or weight only quant ops are
|
||
|
|
run during the convert step.
|
||
|
|
"""
|
||
|
|
for node in observed.graph.nodes:
|
||
|
|
if node.op != "call_function":
|
||
|
|
continue
|
||
|
|
for node_arg in node.args:
|
||
|
|
# node_arg is weight
|
||
|
|
if node_arg and node_arg_is_weight(node, node_arg):
|
||
|
|
weight_observer_nodes = collect_producer_nodes(node_arg)
|
||
|
|
if weight_observer_nodes is None:
|
||
|
|
continue
|
||
|
|
weight_observer_module = \
|
||
|
|
graph_module_from_producer_nodes(
|
||
|
|
observed, weight_observer_nodes)
|
||
|
|
# run the weight observer
|
||
|
|
weight_observer_module()
|
||
|
|
|
||
|
|
def _maybe_recursive_remove_dequantize(arg: Any, node: Node, graph: Graph) -> None:
|
||
|
|
""" If the arg is a dequantize Node, or a list/tuple/dict of dequantize Node,
|
||
|
|
we'll recursively remove the dequantize Node
|
||
|
|
"""
|
||
|
|
if isinstance(arg, Node) and \
|
||
|
|
arg.op == "call_method" and \
|
||
|
|
arg.target == "dequantize":
|
||
|
|
quantize_node = arg.args[0]
|
||
|
|
# we only replace the specific use since dequantize could be used by other nodes
|
||
|
|
# as well
|
||
|
|
node.replace_input_with(arg, quantize_node)
|
||
|
|
elif isinstance(arg, (list, tuple)):
|
||
|
|
for arg_element in arg:
|
||
|
|
_maybe_recursive_remove_dequantize(arg_element, node, graph)
|
||
|
|
elif isinstance(arg, dict):
|
||
|
|
for arg_element in arg.values():
|
||
|
|
_maybe_recursive_remove_dequantize(arg_element, node, graph)
|
||
|
|
else:
|
||
|
|
warnings.warn(f"Unsupported node type in recursive remove dequantize: {type(arg)}")
|
||
|
|
|
||
|
|
def _get_module_path_and_prefix(
|
||
|
|
obs_node: Node,
|
||
|
|
node_name_to_scope: Dict[str, Tuple[str, type]],
|
||
|
|
node_name_to_qconfig: Dict[str, QConfigAny]) -> Tuple[str, str]:
|
||
|
|
""" Given and observer node, get the `Scope` or the fully qualified name for
|
||
|
|
the submodule containing the observed node, also return a prefix of "_input"
|
||
|
|
when the observed node is an input of a F.linear op, and not the output of another
|
||
|
|
quantized op.
|
||
|
|
TODO: this logic is hacky, we should think about how to remove it or make it more
|
||
|
|
general
|
||
|
|
"""
|
||
|
|
observed_node = obs_node.args[0]
|
||
|
|
# an observer can be inserted for both input of the next operator or output of the previous
|
||
|
|
# operator (they can be the same)
|
||
|
|
# this flag identifies if the observer is inserted only because the observed node is
|
||
|
|
# the input of the next operator
|
||
|
|
assert isinstance(observed_node, Node), \
|
||
|
|
f"Expecting observed node to be a Node, but got {observed_node}"
|
||
|
|
is_input_observer_only = node_name_to_qconfig[observed_node.name] is None \
|
||
|
|
if observed_node.name in node_name_to_qconfig else None
|
||
|
|
if is_input_observer_only:
|
||
|
|
# if the quantize function is at the input of op, then we find the first user of the observer_node
|
||
|
|
# to get the path. If a linear call_function is in the user list, we return the first instance
|
||
|
|
# of linear node to get the FQN.
|
||
|
|
users = list(obs_node.users)
|
||
|
|
first_linear_use_or_first_use = users[0] if users else None
|
||
|
|
linear_node = None
|
||
|
|
for n in users:
|
||
|
|
if n.op == "call_function" and n.target == torch.nn.functional.linear:
|
||
|
|
linear_node = n
|
||
|
|
break
|
||
|
|
if linear_node:
|
||
|
|
first_linear_use_or_first_use = linear_node
|
||
|
|
prefix = "_input"
|
||
|
|
else:
|
||
|
|
# if the quantize function is at the output of the op, we use the observer input node to get the path
|
||
|
|
first_linear_use_or_first_use = observed_node
|
||
|
|
prefix = ""
|
||
|
|
|
||
|
|
if first_linear_use_or_first_use and first_linear_use_or_first_use.name in node_name_to_scope:
|
||
|
|
module_path, _ = node_name_to_scope[first_linear_use_or_first_use.name]
|
||
|
|
else:
|
||
|
|
# TODO: it's not used, so actually we can skip quantization
|
||
|
|
# but this requires changing return type of quantize_node
|
||
|
|
# we can fix it later if needed
|
||
|
|
module_path = ""
|
||
|
|
return module_path, prefix
|
||
|
|
|
||
|
|
def _insert_dequantize_node(
|
||
|
|
node: Node,
|
||
|
|
graph: Graph) -> None:
|
||
|
|
""" Inserts dequantize node for `node` in `graph`
|
||
|
|
"""
|
||
|
|
with graph.inserting_after(node):
|
||
|
|
dequantize_node = graph.call_method("dequantize", (node,))
|
||
|
|
for user_node in dict(node.users):
|
||
|
|
if user_node is not dequantize_node:
|
||
|
|
user_node.replace_input_with(node, dequantize_node)
|
||
|
|
|
||
|
|
def _maybe_get_observer_for_node(
|
||
|
|
node: Node,
|
||
|
|
modules: Dict[str, torch.nn.Module]
|
||
|
|
) -> Optional[torch.nn.Module]:
|
||
|
|
"""
|
||
|
|
If the node is observed, return the observer
|
||
|
|
instance. Otherwise, return None.
|
||
|
|
"""
|
||
|
|
for maybe_obs_node in node.users.keys():
|
||
|
|
if maybe_obs_node.op == 'call_module':
|
||
|
|
maybe_obs = modules[str(maybe_obs_node.target)]
|
||
|
|
if _is_activation_post_process(maybe_obs):
|
||
|
|
return maybe_obs
|
||
|
|
return None
|
||
|
|
|
||
|
|
def convert_standalone_module(
|
||
|
|
node: Node,
|
||
|
|
modules: Dict[str, torch.nn.Module],
|
||
|
|
model: torch.fx.GraphModule,
|
||
|
|
is_reference: bool,
|
||
|
|
backend_config: Optional[BackendConfig]) -> None:
|
||
|
|
""" Converts a observed standalone module to a quantized standalone module by calling
|
||
|
|
the fx convert api, currently using the same `is_reference` flag as parent, but we may
|
||
|
|
changing this behavior in the future (e.g. separating quantization and lowering for
|
||
|
|
standalone module as well)
|
||
|
|
|
||
|
|
Args:
|
||
|
|
- node: The call_module node of the observed standalone module
|
||
|
|
- modules: named_module of original model
|
||
|
|
- model: original model
|
||
|
|
- is_reference: a flag from parent provided by user to decide if we want to
|
||
|
|
produce a reference model or a fbgemm/qnnpack model
|
||
|
|
- backend_config: backend configuration of the target backend of quantization
|
||
|
|
"""
|
||
|
|
# TODO: remove is_reference flag
|
||
|
|
if is_reference:
|
||
|
|
convert_fn = torch.ao.quantization.quantize_fx.convert_to_reference_fx
|
||
|
|
else:
|
||
|
|
convert_fn = torch.ao.quantization.quantize_fx.convert_fx # type: ignore[attr-defined]
|
||
|
|
# We know that observed standalone module is a GraphModule since
|
||
|
|
# it's produced by us
|
||
|
|
observed_standalone_module : GraphModule = modules[str(node.target)] # type: ignore[assignment]
|
||
|
|
sm_input_quantized_idxs = \
|
||
|
|
observed_standalone_module \
|
||
|
|
.meta["_observed_graph_module_attrs"].standalone_module_input_quantized_idxs
|
||
|
|
# remove the dequantize nodes for inputs
|
||
|
|
args = list(node.args)
|
||
|
|
for idx in range(len(args)):
|
||
|
|
if idx in sm_input_quantized_idxs:
|
||
|
|
arg = args[idx]
|
||
|
|
if arg.op == "call_method" and arg.target == "dequantize": # type: ignore[union-attr]
|
||
|
|
quantize_node = arg.args[0] # type: ignore[union-attr]
|
||
|
|
node.replace_input_with(arg, quantize_node)
|
||
|
|
if len(arg.users) == 0: # type: ignore[union-attr]
|
||
|
|
model.graph.erase_node(arg)
|
||
|
|
# add dequantize node for output
|
||
|
|
sm_output_quantized_idxs = \
|
||
|
|
observed_standalone_module \
|
||
|
|
.meta["_observed_graph_module_attrs"].standalone_module_output_quantized_idxs
|
||
|
|
if len(sm_output_quantized_idxs) > 0:
|
||
|
|
assert sm_output_quantized_idxs[0] == 0, "Currently only quantized"
|
||
|
|
"output idxs = [0] is supported"
|
||
|
|
|
||
|
|
# if it's non-empty, then it means the output is kept in quantized form
|
||
|
|
# we'll just add a dequantize node after this node
|
||
|
|
_insert_dequantize_node(node, model.graph)
|
||
|
|
|
||
|
|
# TODO: allow convert_custom_config to override backend_config
|
||
|
|
# for standalone module
|
||
|
|
quantized_standalone_module = convert_fn(
|
||
|
|
observed_standalone_module,
|
||
|
|
backend_config=backend_config)
|
||
|
|
parent_name, name = _parent_name(node.target)
|
||
|
|
# update the modules dict
|
||
|
|
setattr(modules[parent_name], name, quantized_standalone_module)
|
||
|
|
modules[str(node.target)] = quantized_standalone_module
|
||
|
|
|
||
|
|
def convert_weighted_module(
|
||
|
|
node: Node,
|
||
|
|
modules: Dict[str, torch.nn.Module],
|
||
|
|
observed_node_names: Set[str],
|
||
|
|
node_name_to_qconfig: Dict[str, QConfigAny],
|
||
|
|
backend_config: BackendConfig,
|
||
|
|
is_decomposed: bool = False,
|
||
|
|
is_reference: bool = False,
|
||
|
|
) -> None:
|
||
|
|
""" Convert a weighted module to reference quantized module in the model
|
||
|
|
If the QConfig of a QAT module is not set, the module will still be converted to
|
||
|
|
a float module.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
- node: The call_module node of the observed standalone module
|
||
|
|
- modules: named_module of original model
|
||
|
|
- observed_node_names: names for the set of observed fx node, we can skip
|
||
|
|
this conversion if the node is not observed
|
||
|
|
"""
|
||
|
|
original_module = modules[str(node.target)]
|
||
|
|
qconfig: QConfigAny = original_module.qconfig # type: ignore[assignment]
|
||
|
|
weight_post_process = None
|
||
|
|
qat_module_classes = get_qat_module_classes(backend_config)
|
||
|
|
|
||
|
|
if isinstance(
|
||
|
|
original_module,
|
||
|
|
qat_module_classes):
|
||
|
|
# Converting qat module to a float module, we need to attach
|
||
|
|
# weight fake_quant to the module, weight fake_quant is assumed to be run during
|
||
|
|
# QAT so we don't need to run it again here
|
||
|
|
weight_post_process = original_module.weight_fake_quant
|
||
|
|
original_module = original_module.to_float() # type: ignore[operator]
|
||
|
|
# change qat module to float module
|
||
|
|
parent_name, name = _parent_name(node.target)
|
||
|
|
setattr(modules[parent_name], name, original_module)
|
||
|
|
|
||
|
|
is_observed = node.name in observed_node_names
|
||
|
|
# If a qconfig is not defined for this node, then skip converting to a reference module
|
||
|
|
if qconfig is None or _has_none_qconfig(node, node_name_to_qconfig) or not is_observed:
|
||
|
|
return
|
||
|
|
|
||
|
|
# skip converting to reference quantized module if the qconfig is not supported
|
||
|
|
pattern_to_dtype_configs = get_pattern_to_dtype_configs(backend_config)
|
||
|
|
dtype_configs = pattern_to_dtype_configs.get(type(original_module), [])
|
||
|
|
if not _is_qconfig_supported_by_dtype_configs(qconfig, dtype_configs):
|
||
|
|
return
|
||
|
|
|
||
|
|
# TODO: rename weight_is_statically_quantized to weight_is_int8_quantized
|
||
|
|
is_weight_quantized = weight_is_quantized(qconfig)
|
||
|
|
|
||
|
|
# the condition for swapping the module to reference quantized module is:
|
||
|
|
# weights need to be quantized
|
||
|
|
if not is_weight_quantized:
|
||
|
|
return
|
||
|
|
|
||
|
|
fused_module = None
|
||
|
|
float_module = original_module
|
||
|
|
# extract the individual float_module and fused module
|
||
|
|
if isinstance(original_module, torch.ao.nn.intrinsic._FusedModule):
|
||
|
|
fused_module = float_module
|
||
|
|
float_module = fused_module[0] # type: ignore[index]
|
||
|
|
|
||
|
|
# TODO: move this to the reference quantized module
|
||
|
|
# weight_qparams or weight_qparams dict
|
||
|
|
wq_or_wq_dict = {"is_decomposed": is_decomposed}
|
||
|
|
if isinstance(float_module, torch.nn.RNNCellBase):
|
||
|
|
weight_post_process_ih = qconfig.weight() # type: ignore[union-attr, operator]
|
||
|
|
weight_post_process_hh = qconfig.weight() # type: ignore[union-attr, operator]
|
||
|
|
weight_post_process_ih(float_module.weight_ih)
|
||
|
|
weight_post_process_hh(float_module.weight_hh)
|
||
|
|
weight_qparams_ih = get_qparam_dict(weight_post_process_ih)
|
||
|
|
weight_qparams_hh = get_qparam_dict(weight_post_process_hh)
|
||
|
|
wq_or_wq_dict.update({
|
||
|
|
"weight_ih": weight_qparams_ih,
|
||
|
|
"weight_hh": weight_qparams_hh,
|
||
|
|
})
|
||
|
|
elif isinstance(float_module, (torch.nn.LSTM, torch.nn.GRU)):
|
||
|
|
# format for wq_or_wq_dict (flattened attributes):
|
||
|
|
# {"weight_ih_l0_scale": ..., "weight_ih_l0_qscheme": ..., ...}
|
||
|
|
for wn in float_module._flat_weights_names:
|
||
|
|
if hasattr(float_module, wn) and wn.startswith("weight"):
|
||
|
|
weight = getattr(float_module, wn)
|
||
|
|
weight_post_process = qconfig.weight() # type: ignore[union-attr, operator]
|
||
|
|
if weight_post_process.dtype == torch.qint8: # type: ignore[union-attr]
|
||
|
|
weight_post_process(weight) # type: ignore[operator, misc]
|
||
|
|
wq_or_wq_dict[wn] = get_qparam_dict(weight_post_process)
|
||
|
|
else:
|
||
|
|
# weight_post_process is None means the original module is not a QAT module
|
||
|
|
# we need to get weight_post_process from qconfig in this case
|
||
|
|
is_ptq = weight_post_process is None
|
||
|
|
if is_ptq:
|
||
|
|
weight_post_process = qconfig.weight() # type: ignore[union-attr, operator]
|
||
|
|
device = assert_and_get_unique_device(float_module)
|
||
|
|
if device:
|
||
|
|
weight_post_process.to(device)
|
||
|
|
|
||
|
|
# Call weight observer/fake_quant at least once to ensure the scales and zero points
|
||
|
|
# have the right shapes. Note: there are two cases where we don't have to do this:
|
||
|
|
#
|
||
|
|
# (1) QAT: The model's forward method already calls the weight observer/fake_quant,
|
||
|
|
# and this typically happens during training, so we don't need to do it here.
|
||
|
|
#
|
||
|
|
# (2) Non-reference (lowered) case: The quantized module's from_float method already
|
||
|
|
# calls the weight observer/fake_quant, so we don't have to do it here.
|
||
|
|
#
|
||
|
|
# Currently we ignore both cases and call the weight observer/fake_quant here
|
||
|
|
# regardless, which is technically incorrect. For (1), this is mainly to preserve BC
|
||
|
|
# in test code, which may not always train before convert. In the future, we should
|
||
|
|
# break BC for these two cases. See https://github.com/pytorch/pytorch/issues/73941.
|
||
|
|
#
|
||
|
|
# For PT2, however, we don't need to preserve BC here, so we can skip this hack
|
||
|
|
# for QAT. We identify this case as (is_decomposed + is_reference + is_qat).
|
||
|
|
# Note that we still need it for PTQ in the PT2 flow since the model's forward
|
||
|
|
# method doesn't call the weight observer.
|
||
|
|
is_qat = not is_ptq
|
||
|
|
if not (is_decomposed and is_reference and is_qat):
|
||
|
|
weight_post_process(float_module.weight) # type: ignore[operator]
|
||
|
|
|
||
|
|
wq_or_wq_dict.update(get_qparam_dict(weight_post_process))
|
||
|
|
|
||
|
|
# We use the same reference module for all modes of quantization: static, dynamic, weight_only
|
||
|
|
# root_module_to_quantized_reference_module: module mapping from root (floating point) module class
|
||
|
|
# to quantized reference module class, e.g. nn.Conv2d to nn.quantized._reference.Conv2d
|
||
|
|
root_module_to_quantized_reference_module = get_root_module_to_quantized_reference_module(backend_config)
|
||
|
|
ref_qmodule_cls = root_module_to_quantized_reference_module.get(type_before_parametrizations(float_module), None)
|
||
|
|
assert (
|
||
|
|
ref_qmodule_cls is not None
|
||
|
|
), f"No reference quantized module class configured for {type_before_parametrizations(float_module)}"
|
||
|
|
ref_qmodule = ref_qmodule_cls.from_float(float_module, wq_or_wq_dict) # type: ignore[attr-defined]
|
||
|
|
if fused_module is not None:
|
||
|
|
fused_module[0] = ref_qmodule # type: ignore[operator]
|
||
|
|
else:
|
||
|
|
parent_name, name = _parent_name(node.target)
|
||
|
|
setattr(modules[parent_name], name, ref_qmodule)
|
||
|
|
|
||
|
|
def _remove_previous_dequantize_in_custom_module(node: Node, prev_node: Node, graph: Graph) -> None:
|
||
|
|
"""
|
||
|
|
Given a custom module `node`, if the previous node is a dequantize, reroute the custom as follows:
|
||
|
|
|
||
|
|
Before: quantize - dequantize - custom_module
|
||
|
|
After: quantize - custom_module
|
||
|
|
\\ - dequantize
|
||
|
|
"""
|
||
|
|
# expecting the input node for a custom module node to be a Node
|
||
|
|
assert isinstance(prev_node, Node), \
|
||
|
|
f"Expecting the argument for custom module node to be a Node, but got {prev_node}"
|
||
|
|
if prev_node.op == "call_method" and prev_node.target == "dequantize":
|
||
|
|
node.replace_input_with(prev_node, prev_node.args[0])
|
||
|
|
# Remove the dequantize node if it doesn't have other users
|
||
|
|
if len(prev_node.users) == 0:
|
||
|
|
graph.erase_node(prev_node)
|
||
|
|
|
||
|
|
def convert_custom_module(
|
||
|
|
node: Node,
|
||
|
|
graph: Graph,
|
||
|
|
modules: Dict[str, torch.nn.Module],
|
||
|
|
custom_module_class_mapping: Dict[QuantType, Dict[Type, Type]],
|
||
|
|
statically_quantized_custom_module_nodes: Set[Node]) -> None:
|
||
|
|
""" Converts an observed custom module to a quantized custom module based on
|
||
|
|
`custom_module_class_mapping`
|
||
|
|
For static quantization, we'll also remove the previous `dequantize` node and
|
||
|
|
attach the observer node for output to the module, the observer for the node
|
||
|
|
will be converted to a dequantize node instead of quantize-dequantize pairs
|
||
|
|
later in the graph. In the end we would have a quantized custom module that
|
||
|
|
has the same interface as a default quantized module in nn.quantized namespace,
|
||
|
|
i.e. quantized input and quantized output.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
- node: The call_module node of the observed standalone module
|
||
|
|
- graph: The graph containing the node
|
||
|
|
- modules: named_module of original model
|
||
|
|
- custom_module_class_mapping: mapping from observed custom module class to
|
||
|
|
quantized custom module class, used to swap custom modules
|
||
|
|
- statically_quantized_custom_module_nodes: we'll add the custom module node
|
||
|
|
if we find it is statically quantized, this will be used later when converting
|
||
|
|
observers to quant/dequant node pairs, if the observed node is a statically
|
||
|
|
quantized custom module nodes, we'll convert the observer to a dequantize node,
|
||
|
|
this is to keep the interface the same as the default quantized module.
|
||
|
|
TODO: maybe we want to redesign this part to align with reference model design
|
||
|
|
as well, but there has been some discussions around the interface, so we can do
|
||
|
|
it later.
|
||
|
|
"""
|
||
|
|
observed_custom_module = modules[str(node.target)]
|
||
|
|
maybe_obs = _maybe_get_observer_for_node(node, modules)
|
||
|
|
qconfig = observed_custom_module.qconfig
|
||
|
|
if activation_is_statically_quantized(qconfig):
|
||
|
|
statically_quantized_custom_module_nodes.add(node)
|
||
|
|
if _is_custom_module_lstm(node, modules):
|
||
|
|
# The inputs are tuples in the form (input, (hidden0, hidden1))
|
||
|
|
# Ensure all three input nodes are quantized
|
||
|
|
assert (
|
||
|
|
len(node.args) == 2 and
|
||
|
|
isinstance(node.args[1], tuple) and
|
||
|
|
len(node.args[1]) == 2
|
||
|
|
)
|
||
|
|
(inputs, (hidden0, hidden1)) = node.args # type: ignore[misc]
|
||
|
|
assert isinstance(inputs, Node)
|
||
|
|
assert isinstance(hidden0, Node)
|
||
|
|
assert isinstance(hidden1, Node)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, inputs, graph)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, hidden0, graph)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, hidden1, graph)
|
||
|
|
elif _is_custom_module_mha(node, modules):
|
||
|
|
# Inputs are in the form (query, key, value)
|
||
|
|
# TODO: This is the first step in enabling the full fx custom module
|
||
|
|
# quantization path for MultiheadAttention, and only covers the inputs
|
||
|
|
# to the module.
|
||
|
|
# Additional handling is yet to be implemented for the outputs, similar
|
||
|
|
# to LSTM custom module
|
||
|
|
assert len(node.args) == 3
|
||
|
|
query, key, value = node.args
|
||
|
|
assert isinstance(query, Node)
|
||
|
|
assert isinstance(key, Node)
|
||
|
|
assert isinstance(value, Node)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, query, graph)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, key, graph)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, value, graph)
|
||
|
|
else:
|
||
|
|
# remove the previous dequant node to ensure the inputs are quantized
|
||
|
|
arg = node.args[0]
|
||
|
|
assert isinstance(arg, Node)
|
||
|
|
_remove_previous_dequantize_in_custom_module(node, arg, graph)
|
||
|
|
# absorb the following observer into the module conversion
|
||
|
|
activation_post_process = _maybe_get_observer_for_node(node, modules)
|
||
|
|
assert activation_post_process is not None
|
||
|
|
observed_custom_module.activation_post_process = activation_post_process
|
||
|
|
|
||
|
|
# swap the observed custom module to quantized custom module
|
||
|
|
quantized_custom_module_class = get_swapped_custom_module_class(
|
||
|
|
observed_custom_module, custom_module_class_mapping, qconfig)
|
||
|
|
quantized_custom_module = \
|
||
|
|
quantized_custom_module_class.from_observed(observed_custom_module)
|
||
|
|
parent_name, name = _parent_name(node.target)
|
||
|
|
setattr(modules[parent_name], name, quantized_custom_module)
|
||
|
|
|
||
|
|
def convert(
|
||
|
|
model: GraphModule, is_reference: bool = False,
|
||
|
|
convert_custom_config: Union[ConvertCustomConfig, Dict[str, Any], None] = None,
|
||
|
|
is_standalone_module: bool = False,
|
||
|
|
_remove_qconfig_flag: bool = True,
|
||
|
|
qconfig_mapping: Union[QConfigMapping, Dict[str, Any], None] = None,
|
||
|
|
backend_config: Union[BackendConfig, Dict[str, Any], None] = None,
|
||
|
|
is_decomposed: bool = False) -> GraphModule:
|
||
|
|
"""
|
||
|
|
We will convert an observed model (a module with observer calls) to a reference
|
||
|
|
quantized model, the rule is simple:
|
||
|
|
1. for each observer module call in the graph, we'll convert it to calls to
|
||
|
|
quantize and dequantize functions based on the observer instance
|
||
|
|
2. for weighted operations like linear/conv, we need to convert them to reference
|
||
|
|
quantized module, this requires us to know whether the dtype configured for the
|
||
|
|
weight is supported in the backend, this is done in prepare step and the result
|
||
|
|
is stored in observed_node_names, we can decide whether we need to swap the
|
||
|
|
module based on this set
|
||
|
|
|
||
|
|
Args:
|
||
|
|
* `is_standalone_module`: when this flag is True, it means we are quantizing
|
||
|
|
a submodule that is not inlined in parent module, and will be quantized
|
||
|
|
separately as one unit.
|
||
|
|
|
||
|
|
* `is_decomposed`: a boolean flag to indicate whether we want to use the
|
||
|
|
quantize operator for decomposed quantized tensor
|
||
|
|
(torch.ops.quantized_decomposed.quantize_per_tensor) or default/standalone
|
||
|
|
quantized tensor (torch.quantize_per_tensor)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
a quantized standalone module, whether input/output is quantized is
|
||
|
|
specified by prepare_custom_config, with
|
||
|
|
input_quantized_idxs, output_quantized_idxs, please
|
||
|
|
see docs for :func:`~torch.ao.quantization.prepare_fx` for details
|
||
|
|
"""
|
||
|
|
if convert_custom_config is None:
|
||
|
|
convert_custom_config = ConvertCustomConfig()
|
||
|
|
|
||
|
|
if isinstance(convert_custom_config, Dict):
|
||
|
|
warnings.warn(
|
||
|
|
"Passing a convert_custom_config_dict to convert is deprecated and will not be supported "
|
||
|
|
"in a future version. Please pass in a ConvertCustomConfig instead.")
|
||
|
|
convert_custom_config = ConvertCustomConfig.from_dict(convert_custom_config)
|
||
|
|
|
||
|
|
if isinstance(qconfig_mapping, Dict):
|
||
|
|
warnings.warn(
|
||
|
|
"Passing a QConfig dictionary to convert is deprecated and will not be supported "
|
||
|
|
"in a future version. Please pass in a QConfigMapping instead.")
|
||
|
|
qconfig_mapping = QConfigMapping.from_dict(qconfig_mapping) if qconfig_mapping else None
|
||
|
|
qconfig_mapping = copy.deepcopy(qconfig_mapping)
|
||
|
|
assert qconfig_mapping is None or isinstance(qconfig_mapping, QConfigMapping)
|
||
|
|
|
||
|
|
if isinstance(backend_config, Dict):
|
||
|
|
warnings.warn(
|
||
|
|
"Passing a backend_config_dict to prepare is deprecated and will not be supported "
|
||
|
|
"in a future version. Please pass in a BackendConfig instead.")
|
||
|
|
backend_config = BackendConfig.from_dict(backend_config)
|
||
|
|
|
||
|
|
if backend_config is None:
|
||
|
|
backend_config = get_native_backend_config()
|
||
|
|
|
||
|
|
assert _is_observed_module(model), \
|
||
|
|
'incoming model must be produced by prepare_fx'
|
||
|
|
observed_graph_module_attrs = model.meta["_observed_graph_module_attrs"]
|
||
|
|
node_name_to_scope: Dict[str, Tuple[str, type]] = observed_graph_module_attrs.node_name_to_scope
|
||
|
|
prepare_custom_config: PrepareCustomConfig = observed_graph_module_attrs.prepare_custom_config
|
||
|
|
observed_node_names: Set[str] = observed_graph_module_attrs.observed_node_names
|
||
|
|
node_name_to_qconfig: Dict[str, QConfigAny] = observed_graph_module_attrs.node_name_to_qconfig # type: ignore[assignment]
|
||
|
|
|
||
|
|
# mapping from fully qualified module name to module instance
|
||
|
|
# for example,
|
||
|
|
# {
|
||
|
|
# '': Model(...),
|
||
|
|
# 'linear': Linear(...),
|
||
|
|
# 'linear.weight_fake_quant': PerChannelMinMaxObserver(...),
|
||
|
|
# }
|
||
|
|
# We use remove_duplicate=False here because torch.cat uses
|
||
|
|
# the same activation_post_process module instance but different names
|
||
|
|
modules = dict(model.named_modules(remove_duplicate=False))
|
||
|
|
|
||
|
|
# TODO refactor this code once we update the prepare logic to have additional information on
|
||
|
|
# which graph nodes have been observed and share that with convert to decide which observers to ignore.
|
||
|
|
if qconfig_mapping:
|
||
|
|
prepare_qconfig_mapping: QConfigMapping = observed_graph_module_attrs.qconfig_mapping # type: ignore[assignment]
|
||
|
|
modules_copy = copy.deepcopy(modules)
|
||
|
|
|
||
|
|
if observed_graph_module_attrs.is_qat:
|
||
|
|
_update_qconfig_for_qat(qconfig_mapping, backend_config)
|
||
|
|
_update_qconfig_for_fusion(model, qconfig_mapping)
|
||
|
|
|
||
|
|
_compare_prepare_convert_qconfig_mappings(prepare_qconfig_mapping, qconfig_mapping) # type: ignore[arg-type]
|
||
|
|
convert_node_name_to_qconfig = _generate_node_name_to_qconfig(
|
||
|
|
model, modules_copy, model.graph, qconfig_mapping, node_name_to_scope)
|
||
|
|
# check the convert_node_name_to_qconfig generated and ensure that
|
||
|
|
# all the values either match what was set in prepare node_name_to_qconfig
|
||
|
|
# or are set to None in the convert_node_name_to_qconfig.
|
||
|
|
for k, v in node_name_to_qconfig.items():
|
||
|
|
assert k in convert_node_name_to_qconfig, f'Expected key {k} in convert node_name_to_qconfig'
|
||
|
|
if convert_node_name_to_qconfig[k] is not None:
|
||
|
|
assert qconfig_equals(v, convert_node_name_to_qconfig[k]), \
|
||
|
|
f"Expected k {k} to have the same value in prepare and convert QConfigMappings, " \
|
||
|
|
f"but {v} was updated to {convert_node_name_to_qconfig[k]}"
|
||
|
|
node_name_to_qconfig = convert_node_name_to_qconfig
|
||
|
|
|
||
|
|
custom_module_classes = get_custom_module_class_keys(convert_custom_config.observed_to_quantized_mapping)
|
||
|
|
custom_module_class_mapping = convert_custom_config.observed_to_quantized_mapping
|
||
|
|
|
||
|
|
if observed_graph_module_attrs.equalization_node_name_to_qconfig is not None:
|
||
|
|
# If we want to do equalization then do the following:
|
||
|
|
# Calculate the equalization scale, update the observers with the scaled
|
||
|
|
# inputs, and scale the weight
|
||
|
|
weight_eq_obs_dict = update_obs_for_equalization(model, modules)
|
||
|
|
convert_eq_obs(model, modules, weight_eq_obs_dict)
|
||
|
|
|
||
|
|
# always run weight observers in the top level forward method
|
||
|
|
# for dynamic quant ops or weight only quant ops
|
||
|
|
_run_weight_observers(model, backend_config)
|
||
|
|
|
||
|
|
graph_inputs: List[str] = []
|
||
|
|
for node in model.graph.nodes:
|
||
|
|
if node.op == 'placeholder':
|
||
|
|
graph_inputs.append(node.name)
|
||
|
|
|
||
|
|
# additional state to override inputs to be quantized, if specified
|
||
|
|
# by the user
|
||
|
|
placeholder_node_seen_cnt = 0
|
||
|
|
input_quantized_idxs: List[int] = prepare_custom_config.input_quantized_indexes
|
||
|
|
output_quantized_idxs: List[int] = prepare_custom_config.output_quantized_indexes
|
||
|
|
|
||
|
|
root_module_to_quantized_reference_module = get_root_module_to_quantized_reference_module(backend_config)
|
||
|
|
# convert tuples so that it can work with isinstance(module, tuple_of_classes)
|
||
|
|
root_module_classes = tuple(root_module_to_quantized_reference_module.keys())
|
||
|
|
qat_module_classes = get_qat_module_classes(backend_config)
|
||
|
|
fused_module_classes = get_fused_module_classes(backend_config)
|
||
|
|
statically_quantized_custom_module_nodes: Set[Node] = set()
|
||
|
|
|
||
|
|
for node in list(model.graph.nodes):
|
||
|
|
if node.op == 'placeholder':
|
||
|
|
cur_placeholder_node_idx = placeholder_node_seen_cnt
|
||
|
|
placeholder_node_seen_cnt += 1
|
||
|
|
if cur_placeholder_node_idx in input_quantized_idxs:
|
||
|
|
# Inputs are assumed to be quantized if the user specified the
|
||
|
|
# input_quantized_idxs override.
|
||
|
|
# we need to dequantize the inputs since all operators took
|
||
|
|
# floating point inputs in reference quantized models
|
||
|
|
_insert_dequantize_node(node, model.graph)
|
||
|
|
elif node.op == "output":
|
||
|
|
# If the argument is empty we don't need to do anything
|
||
|
|
if len(output_quantized_idxs) == 0:
|
||
|
|
continue
|
||
|
|
# Result are kept quantized if the user specified the
|
||
|
|
# output_quantized_idxs override.
|
||
|
|
# Remove the dequantize operator for the node in the end if any
|
||
|
|
return_node = node
|
||
|
|
output = node.args[0]
|
||
|
|
# outputs can be Node, list, tuple, dict, other cases are not supported yet
|
||
|
|
if isinstance(output, (list, tuple)):
|
||
|
|
for idx in output_quantized_idxs:
|
||
|
|
_maybe_recursive_remove_dequantize(output[idx], return_node, model.graph)
|
||
|
|
elif isinstance(output, (Node, dict)):
|
||
|
|
# we treat dict as a single argument currently, but it can be extended
|
||
|
|
# to support {"key": dtype} after we change output_quantized_idxs to
|
||
|
|
# dict
|
||
|
|
if 0 in output_quantized_idxs:
|
||
|
|
_maybe_recursive_remove_dequantize(output, return_node, model.graph)
|
||
|
|
else:
|
||
|
|
warnings.warn(f"Unsupported node type for output_quantized_idxs: {type(output)}")
|
||
|
|
elif node.op == "call_module":
|
||
|
|
mod = _get_module(node, modules)
|
||
|
|
assert mod is not None
|
||
|
|
if _is_activation_post_process(mod):
|
||
|
|
observed_node = node.args[0]
|
||
|
|
if observed_node in statically_quantized_custom_module_nodes:
|
||
|
|
_replace_observer_or_dequant_stub_with_dequantize_node(node, model.graph)
|
||
|
|
else:
|
||
|
|
if is_decomposed:
|
||
|
|
_replace_observer_with_quantize_dequantize_node_decomposed(
|
||
|
|
model, node, modules, node_name_to_scope,
|
||
|
|
node_name_to_qconfig)
|
||
|
|
else:
|
||
|
|
_replace_observer_with_quantize_dequantize_node(
|
||
|
|
model, node, modules, node_name_to_scope,
|
||
|
|
node_name_to_qconfig)
|
||
|
|
elif isinstance(mod, DeQuantStub):
|
||
|
|
_replace_observer_or_dequant_stub_with_dequantize_node(node, model.graph)
|
||
|
|
elif _is_observed_standalone_module(mod):
|
||
|
|
convert_standalone_module(
|
||
|
|
node, modules, model, is_reference, backend_config)
|
||
|
|
# below this point `type_before_parametrizations` is used
|
||
|
|
# instead of `type` to handle situations with fx quant + sparsity
|
||
|
|
elif type_before_parametrizations(mod) in set(
|
||
|
|
root_module_classes).union(qat_module_classes).union(fused_module_classes):
|
||
|
|
# extra check for fused module classes to make sure they are fused module classes
|
||
|
|
# of target modules
|
||
|
|
if type_before_parametrizations(mod) in fused_module_classes and \
|
||
|
|
type_before_parametrizations(mod[0]) not in root_module_classes: # type: ignore[index]
|
||
|
|
continue
|
||
|
|
convert_weighted_module(
|
||
|
|
node, modules, observed_node_names, node_name_to_qconfig, backend_config,
|
||
|
|
is_decomposed, is_reference)
|
||
|
|
elif type_before_parametrizations(mod) in custom_module_classes:
|
||
|
|
convert_custom_module(
|
||
|
|
node, model.graph, modules, custom_module_class_mapping,
|
||
|
|
statically_quantized_custom_module_nodes)
|
||
|
|
|
||
|
|
# remove deadcode after converting observers to quant/dequant ops
|
||
|
|
model.graph.eliminate_dead_code()
|
||
|
|
model = GraphModule(model, model.graph)
|
||
|
|
|
||
|
|
# TODO: maybe move this to quantize_fx.py
|
||
|
|
if not is_reference:
|
||
|
|
model = lower_to_fbgemm(model, node_name_to_qconfig, node_name_to_scope)
|
||
|
|
|
||
|
|
# TODO: this looks hacky, we want to check why we need this and see if we can
|
||
|
|
# remove this
|
||
|
|
# removes qconfig and activation_post_process modules
|
||
|
|
if _remove_qconfig_flag:
|
||
|
|
_remove_qconfig(model)
|
||
|
|
model.delete_all_unused_submodules()
|
||
|
|
model.meta.pop("_observed_graph_module_attrs", None)
|
||
|
|
return model
|