diff --git a/orttraining/orttraining/python/training/onnxblock/_graph_utils.py b/orttraining/orttraining/python/training/onnxblock/_graph_utils.py index 0d10ffc0ad..d8dc0fa733 100644 --- a/orttraining/orttraining/python/training/onnxblock/_graph_utils.py +++ b/orttraining/orttraining/python/training/onnxblock/_graph_utils.py @@ -8,6 +8,7 @@ import random import onnx from onnxruntime.capi._pybind_state import GradientGraphBuilder, get_optimized_model +import onnxruntime.training.onnxblock._qat_utils as qat_utils def get_output_from_output_name(onnx_model, output_name): @@ -50,6 +51,8 @@ def build_gradient_graph(accessor, user_args_requiring_grad, user_args_not_requi model = accessor.model + quant_params = qat_utils.get_quant_params(model) + # Collect names of parameters that need gradients computed all_args_requiring_gradient = [] # Move all trainable and non trainable initializers to graph inputs. @@ -58,7 +61,7 @@ def build_gradient_graph(accessor, user_args_requiring_grad, user_args_not_requi graph_inputs = model.graph.input initializers = [] for initializer in model.graph.initializer: - if not initializer.name.startswith("onnx::"): + if not initializer.name.startswith("onnx::") and initializer.name not in quant_params: # Move only those initializers as inputs that are not local # to the onnx model. i.e. initializers that are model parameters. # These are tpically those initializers without any onnx:: prefixed @@ -69,7 +72,7 @@ def build_gradient_graph(accessor, user_args_requiring_grad, user_args_not_requi if initializer.name not in user_args_not_requiring_grad: all_args_requiring_gradient.append(initializer.name) else: - # All other initializers stay where they were. + # All other initializers (including any quantization parameter) stay where they were. initializers.append(initializer) # Update the initializers in the graph @@ -189,6 +192,8 @@ def build_gradient_accumulation_graph(grad_model, all_args_requiring_gradient_na def get_model_parameters(model, args_not_requiring_gradient): """Returns trainable and non trainable onnx model parameters.""" + quant_params = qat_utils.get_quant_params(model) + trainable_params = [] non_trainable_params = [] for initializer in model.graph.initializer: @@ -204,7 +209,9 @@ def get_model_parameters(model, args_not_requiring_gradient): # logic and have a `onnx::` prefix. # TODO: validate this assumption. If assumption is not valid, # the alternative is to enforce the user to provide the parameter names. - if not initializer.name.startswith("onnx::"): + # Do not move quantization parameters to graph inputs, so skip putting them in either + # trainable or non-trainable parameters. + if not initializer.name.startswith("onnx::") and initializer.name not in quant_params: if initializer.name in args_not_requiring_gradient: non_trainable_params.append(initializer) else: diff --git a/orttraining/orttraining/python/training/onnxblock/_qat_utils.py b/orttraining/orttraining/python/training/onnxblock/_qat_utils.py new file mode 100644 index 0000000000..251851610c --- /dev/null +++ b/orttraining/orttraining/python/training/onnxblock/_qat_utils.py @@ -0,0 +1,12 @@ +def get_quant_params(model): + """Returns quantization parameters for the given model. + + Quantization parameters in this function refers to all scale and zero point + inputs to any QuantizeLinear, DequantizeLinear or FakeQuant node in the model.""" + + return { + quant_param_name + for node in model.graph.node + for quant_param_name in node.input[1:] + if node.op_type == "QuantizeLinear" or node.op_type == "DequantizeLinear" or node.op_type == "FakeQuant" + }