diff --git a/paddle/fluid/operators/quantize_linear_op.cc b/paddle/fluid/operators/quantize_linear_op.cc index cab16683b4ac1c18c946957942c63dcb0643d070..7f9d472cb5284fed9f1d09086da8020af0fb45d7 100644 --- a/paddle/fluid/operators/quantize_linear_op.cc +++ b/paddle/fluid/operators/quantize_linear_op.cc @@ -24,6 +24,22 @@ limitations under the License. */ namespace paddle { namespace operators { +template +struct DequantizeFunctor { + void operator()(const phi::CPUContext &dev_ctx, + const phi::DenseTensor *in, + const phi::DenseTensor *scale, + T max_range, + phi::DenseTensor *out) { + auto in_e = framework::EigenVector::Flatten(*in); + const T *scale_factor = scale->data(); + auto out_e = framework::EigenVector::Flatten(*out); + + auto &dev = *dev_ctx.eigen_device(); + out_e.device(dev) = in_e * scale_factor[0] / max_range; + } +}; + template struct ChannelDequantizeFunctorV2 { void operator()(const phi::CPUContext &dev_ctx, @@ -55,7 +71,7 @@ struct ChannelDequantizeFunctorV2 { int64_t step_i = in->numel() / out_iter; int64_t step_j = in->numel() / (out_iter * channel); auto *in_data = in->data(); - auto *out_data = out->mutable_data(dev_ctx.GetPlace()); + auto *out_data = dev_ctx.Alloc(out, out->numel() * sizeof(T)); for (int64_t i = 0; i < out_iter; i++) { for (int64_t j = 0; j < channel; j++) { auto *cur_in = in_data + i * step_i + j * step_j; @@ -72,6 +88,11 @@ struct ChannelDequantizeFunctorV2 { } }; +template struct DequantizeFunctor; +template struct DequantizeFunctor; +template struct DequantizeFunctor; +template struct ChannelDequantizeFunctorV2; template struct ChannelDequantizeFunctorV2; template struct ChannelDequantizeFunctorV2; @@ -214,6 +235,6 @@ REGISTER_OPERATOR( paddle::framework::EmptyGradOpMaker); REGISTER_OP_CPU_KERNEL(dequantize_linear, - ops::DeQuantizeLinearKernel, - ops::DeQuantizeLinearKernel, - ops::DeQuantizeLinearKernel); + ops::DeQuantizeLinearKernel, + ops::DeQuantizeLinearKernel, + ops::DeQuantizeLinearKernel); diff --git a/paddle/fluid/operators/quantize_linear_op.cu b/paddle/fluid/operators/quantize_linear_op.cu index 259c1507af038a0032ffa13ced5468afe9eb779a..92faa08bbbce1104d9e44db1f940ba547add2f31 100644 --- a/paddle/fluid/operators/quantize_linear_op.cu +++ b/paddle/fluid/operators/quantize_linear_op.cu @@ -15,14 +15,64 @@ limitations under the License. */ #include #include "paddle/fluid/memory/memcpy.h" -#include "paddle/fluid/operators/fake_dequantize_op.cu.h" #include "paddle/fluid/operators/fake_quantize_op.cu.h" #include "paddle/fluid/operators/quantize_linear_op.h" #include "paddle/phi/backends/gpu/gpu_primitives.h" +using float16 = paddle::platform::float16; + namespace paddle { namespace operators { +template +__global__ void KeDequantize( + const T* in, const T* scale, T max_range, int64_t num, T* out) { + int64_t idx = threadIdx.x + blockIdx.x * blockDim.x; + for (int64_t i = idx; i < num; i += blockDim.x * gridDim.x) { + out[i] = in[i] * scale[0] / max_range; + } +} + +template +__global__ void DequantizeOneScaleQuantAxisN(const T* in, + const T* scale, + const T max_range, + const int64_t num, + const int n_scales, + const int quant_stride, + T* out) { + int64_t idx = blockDim.x * blockIdx.x + threadIdx.x; + for (int64_t i = idx; i < num; i += blockDim.x * gridDim.x) { + T s = scale[(i / quant_stride) % n_scales]; + out[i] = in[i] * s / max_range; + } +} + +template +struct DequantizeFunctor { + void operator()(const phi::GPUContext& dev_ctx, + const phi::DenseTensor* in, + const phi::DenseTensor* scale, + T max_range, + phi::DenseTensor* out) { + const T* in_data = in->data(); + const T* scale_factor = scale->data(); + T* out_data = dev_ctx.Alloc(out, out->numel() * sizeof(T)); + + int64_t num = in->numel(); + int64_t block_size = std::min( + num, static_cast(dev_ctx.GetMaxThreadsPerBlock() / 4)); + int64_t max_threads = + dev_ctx.GetMaxPhysicalThreadCount(); // SM * block_per_SM + const int64_t max_blocks = + std::max(((max_threads - 1) / block_size + 1), static_cast(1)); + const int64_t grid_size = + std::min(max_blocks, (num + block_size - 1) / block_size); + KeDequantize<<>>( + in_data, scale_factor, max_range, num, out_data); + } +}; + template struct ChannelDequantizeFunctorV2 { void operator()(const phi::GPUContext& dev_ctx, @@ -33,7 +83,7 @@ struct ChannelDequantizeFunctorV2 { phi::DenseTensor* out) { auto in_dims = in->dims(); const T* in_data = in->data(); - T* out_data = out->mutable_data(dev_ctx.GetPlace()); + T* out_data = dev_ctx.Alloc(out, out->numel() * sizeof(T)); int64_t num = in->numel(); const T* scale_factor = scale->data(); int64_t block_size = std::min( @@ -61,6 +111,10 @@ struct ChannelDequantizeFunctorV2 { } }; +template struct DequantizeFunctor; +template struct DequantizeFunctor; +template struct DequantizeFunctor; +template struct ChannelDequantizeFunctorV2; template struct ChannelDequantizeFunctorV2; template struct ChannelDequantizeFunctorV2; @@ -70,9 +124,11 @@ template struct ChannelDequantizeFunctorV2; namespace ops = paddle::operators; using CUDA = phi::GPUContext; REGISTER_OP_CUDA_KERNEL(dequantize_linear, - ops::DeQuantizeLinearKernel, - ops::DeQuantizeLinearKernel, - ops::DeQuantizeLinearKernel); + ops::DeQuantizeLinearKernel, + ops::DeQuantizeLinearKernel, + ops::DeQuantizeLinearKernel, + ops::DeQuantizeLinearKernel); REGISTER_OP_CUDA_KERNEL(quantize_linear, - ops::QuantizeLinearKernel); + ops::QuantizeLinearKernel, + ops::QuantizeLinearKernel); diff --git a/paddle/fluid/operators/quantize_linear_op.h b/paddle/fluid/operators/quantize_linear_op.h index 8434996926aba5a5c4c7e5163c9291686de9c305..3461b9de0a53ae31b6973bcc6e3de9fe1e321e18 100644 --- a/paddle/fluid/operators/quantize_linear_op.h +++ b/paddle/fluid/operators/quantize_linear_op.h @@ -17,7 +17,6 @@ limitations under the License. */ #include "paddle/fluid/framework/op_registry.h" #include "paddle/fluid/framework/tensor_util.h" #include "paddle/fluid/memory/malloc.h" -#include "paddle/fluid/operators/fake_dequantize_op.h" #include "paddle/fluid/operators/fake_quantize_op.h" #include "paddle/fluid/platform/transform.h" #include "paddle/phi/common/data_type.h" @@ -28,6 +27,15 @@ limitations under the License. */ namespace paddle { namespace operators { +template +struct DequantizeFunctor { + void operator()(const DeviceContext& dev_ctx, + const phi::DenseTensor* in, + const phi::DenseTensor* scale, + T max_range, + phi::DenseTensor* out); +}; + template struct ChannelDequantizeFunctorV2 { void operator()(const DeviceContext& dev_ctx, @@ -105,10 +113,11 @@ class QuantizeLinearKernel : public framework::OpKernel { } }; -template +template class DeQuantizeLinearKernel : public framework::OpKernel { public: - void Compute(const framework::ExecutionContext& context) const override { + template + void ComputeImpl(const framework::ExecutionContext& context) const { auto& dev_ctx = context.template device_context(); auto* in = context.Input("X"); @@ -122,7 +131,7 @@ class DeQuantizeLinearKernel : public framework::OpKernel { auto* out = context.Output("Y"); int bit_length = context.Attr("bit_length"); auto quant_axis = context.Attr("quant_axis"); - out->mutable_data(dev_ctx.GetPlace()); + dev_ctx.template Alloc(out, out->numel() * sizeof(D)); if (quant_axis < 0) { float max_range = (std::pow(2, bit_length - 1) - 1); @@ -144,6 +153,27 @@ class DeQuantizeLinearKernel : public framework::OpKernel { dev_ctx, &in_tmp, scale, static_cast(max_range), quant_axis, out); } } + + void Compute(const framework::ExecutionContext& context) const override { + auto* scale = context.Input("Scale"); + switch (scale->dtype()) { + case experimental::DataType::FLOAT64: + ComputeImpl(context); + break; + case experimental::DataType::FLOAT32: + ComputeImpl(context); + break; + case experimental::DataType::FLOAT16: + ComputeImpl(context); + break; + default: + PADDLE_THROW(platform::errors::Unimplemented( + "In DeQuantizeLinearKernel, " + "data type %d for scale/output is not supported ", + scale->dtype())); + break; + } + } }; } // namespace operators diff --git a/python/paddle/distributed/auto_parallel/constants.py b/python/paddle/distributed/auto_parallel/constants.py index ce72304dc75cdf732bf0362c6a82d08ec134155b..044bc78887b03ca0cb4f115385476291625b42b3 100644 --- a/python/paddle/distributed/auto_parallel/constants.py +++ b/python/paddle/distributed/auto_parallel/constants.py @@ -107,6 +107,7 @@ set_field_default_config(QAT, "weight_bits", 8) set_field_default_config(QAT, "activation_bits", 8) set_field_default_config(QAT, "not_quant_pattern", ['skip_quant']) set_field_default_config(QAT, "algo", None) +set_field_default_config(QAT, "onnx_format", True) # ######################################### # auto tuning configuration diff --git a/python/paddle/distributed/auto_parallel/engine.py b/python/paddle/distributed/auto_parallel/engine.py index 171eab28da1f18c086359435d2c06c05a44cfce4..1ef44bdb7f5b723a527008e7fd66ce9b6474701a 100644 --- a/python/paddle/distributed/auto_parallel/engine.py +++ b/python/paddle/distributed/auto_parallel/engine.py @@ -29,7 +29,7 @@ from paddle.distributed import fleet from paddle.fluid import Variable, core from paddle.fluid.dygraph.parallel import ParallelEnv from paddle.fluid.executor import _to_name_str, global_scope -from paddle.fluid.framework import Operator +from paddle.fluid.framework import IrGraph, Operator from paddle.fluid.framework import _current_expected_place as _get_device from paddle.fluid.framework import in_dygraph_mode from paddle.fluid.layers.utils import flatten @@ -752,7 +752,9 @@ class Engine: # instantiate communication by process_mapping. all_process_groups = get_all_process_groups() cur_rank = self._cur_rank - # NOTE: After the implementation of the unified dynamic and static communication group initialization mode in the future, the initialization logic of full mode will be removed because port occupation error may occur. + # NOTE: After the implementation of the unified dynamic and static communication group + # initialization mode in the future, the initialization logic of full mode + # will be removed because port occupation error may occur. if self._strategy.auto_mode == "full": auto_utils.initialize_pg_in_full_mode( all_process_groups, cur_rank @@ -763,9 +765,9 @@ class Engine: continue process_group.instantiate() - place = _get_device() - if isinstance(place, fluid.CUDAPlace): - place = fluid.CUDAPlace(ParallelEnv().dev_id) + self._place = _get_device() + if isinstance(self._place, fluid.CUDAPlace): + self._place = fluid.CUDAPlace(ParallelEnv().dev_id) if self._strategy.seed: paddle.seed(self._strategy.seed + self._dp_ranks[0]) @@ -775,10 +777,12 @@ class Engine: if self._dygraph_mode: dist_context = self._dist_contexts[mode] dist_main_program = self._dist_main_progs[mode][self._cur_rank] - self.program_helper.init(dist_main_program, place, dist_context) + self.program_helper.init( + dist_main_program, self._place, dist_context + ) if self._executor is None: - self._executor = paddle.static.Executor(place) + self._executor = paddle.static.Executor(self._place) uninitialized = [] dist_startup_prog = self._dist_startup_progs[mode][self._cur_rank] for var in dist_startup_prog.list_vars(): @@ -1612,6 +1616,22 @@ class Engine: feed_vars = self._feed_vars["predict"]['inputs'] fetch_vars = self._fetch_vars["predict"]['outputs'] dist_main_prog = self._dist_main_progs["predict"][self._cur_rank] + if self._strategy.qat.enable and self._strategy.qat.onnx_format: + from paddle.fluid.contrib.slim.quantization import ( + QuantWeightPass, + ) + + self._logger.info("export quantized model.") + self._logger.info( + "convert config {}".format(self._strategy.qat.to_dict()) + ) + test_graph = IrGraph( + core.Graph(dist_main_prog.desc), for_test=True + ) + quant_weight_pass = QuantWeightPass(global_scope(), self._place) + for sub_graph in test_graph.all_sub_graphs(): + quant_weight_pass.apply(sub_graph) + dist_main_prog = test_graph.to_program() self._saver.save_inference_model( path, feed_vars, diff --git a/python/paddle/distributed/auto_parallel/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/parallelizer_v2.py index dba27588c7377ea870d49fb481f29f8c385fea90..0a05c716ded64595c8f436b115af5c2fb360ce26 100644 --- a/python/paddle/distributed/auto_parallel/parallelizer_v2.py +++ b/python/paddle/distributed/auto_parallel/parallelizer_v2.py @@ -131,8 +131,12 @@ class Parallelizer: else: # Apply pre optimization passes time0 = time.time() - self._apply_pre_optimization( - serial_main_program, serial_startup_program, None, None, None + ( + serial_main_program, + serial_startup_program, + params_grads, + ) = self._apply_pre_optimization( + serial_main_program, serial_startup_program, None, None, [] ) self._logger.debug( "within parallel apply_pre_optimization time: {}, mode {}".format( @@ -207,22 +211,6 @@ class Parallelizer: if self._strategy is None: return - # apply quantization pass - # The pass can be applied when mode must be 'train' - if self._mode == 'train' and self._strategy.qat.enable: - config = copy.deepcopy(self._strategy.qat.to_dict()) - config["dist_context"] = self._dist_context - config["params_grads"] = params_grads - auto_parallel_quantization_pass = new_pass( - "auto_parallel_quantization", config - ) - auto_parallel_quantization_pass.apply( - [main_program], [startup_program], self._pass_context - ) - main_program = self._pass_context.get_attr("main_program") - startup_program = self._pass_context.get_attr("startup_program") - params_grads = self._pass_context.get_attr("params_grads") - # apply amp pass on train/eval/predict if self._strategy.amp.enable: config = copy.deepcopy(self._strategy.amp.to_dict()) @@ -247,6 +235,25 @@ class Parallelizer: ) loss = auto_parallel_amp_pass.get_loss() + # apply quantization pass + # The pass can be applied when mode must be 'train' + if self._strategy.qat.enable: + config = copy.deepcopy(self._strategy.qat.to_dict()) + config["dist_context"] = self._dist_context + config["params_grads"] = params_grads + config["mode"] = self._mode + config["loss"] = loss + auto_parallel_quantization_pass = new_pass( + "auto_parallel_quantization", config + ) + auto_parallel_quantization_pass.apply( + [main_program], [startup_program], self._pass_context + ) + main_program = self._pass_context.get_attr("main_program") + startup_program = self._pass_context.get_attr("startup_program") + params_grads = self._pass_context.get_attr("params_grads") + loss = self._pass_context.get_attr("loss") + # apply recompute pass # recompute is then train-only optimization if self._mode == "train" and self._strategy.recompute.enable: diff --git a/python/paddle/distributed/auto_parallel/reshard.py b/python/paddle/distributed/auto_parallel/reshard.py index 3dda3fb3899d63903f8ba88f967ee99872088143..07c40828bde18eb4752960889620e7c574ce2b7d 100644 --- a/python/paddle/distributed/auto_parallel/reshard.py +++ b/python/paddle/distributed/auto_parallel/reshard.py @@ -2137,6 +2137,36 @@ class Resharder: input_attrs.append([process_mesh, input_dims_mapping]) return input_attrs + def _get_subblock_output_attrs(self, op, var_name): + # NOTE: Multi while loop is not supported + assert op.type in _g_subblock_ops + sub_block = self.auto_parallel_main_prog.blocks[op.attr("sub_block").id] + ops = sub_block.ops + output_attrs = [] + + for op in ops: + dist_op = self.dist_context.get_dist_op_for_program(op) + if not dist_op: + continue + dist_attr = dist_op.dist_attr + for name in op.output_arg_names: + if name == var_name: + process_mesh = dist_attr.process_mesh + output_dims_mapping = dist_attr.get_output_dims_mapping( + var_name + ) + has_exist = False + for output_attr in output_attrs: + if ( + process_mesh == output_attrs[0] + and output_dims_mapping == output_attrs[1] + ): + has_exist = True + break + if not has_exist: + output_attrs.append([process_mesh, output_dims_mapping]) + return output_attrs + def _get_common_op_input_attrs(self, op, var_name): process_meshes = [] dist_op = self.dist_context.get_dist_op_for_program(op) @@ -2166,6 +2196,11 @@ class Resharder: if op.type in _g_subblock_ops: op_input_attrs = self._get_subblock_input_attrs(op, var_name) + if not op_input_attrs: + # NOTE: [hack method] + # Adapt to quantization pass, which presist_vars, including inputs and outputs, all are in global_block. + # Therefore, the while_op's inputs will contain the all persist_vars, which will be inputs or output of the quantization op in subblock. + op_input_attrs = self._get_subblock_output_attrs(op, var_name) else: op_input_attrs = self._get_common_op_input_attrs(op, var_name) diff --git a/python/paddle/distributed/passes/auto_parallel_amp.py b/python/paddle/distributed/passes/auto_parallel_amp.py index 06b8e4a19e88f90cf0f44b0a89d7237e346f12bb..fd5ddf894716e7ecbe3b6c17d39f899c0e09d915 100644 --- a/python/paddle/distributed/passes/auto_parallel_amp.py +++ b/python/paddle/distributed/passes/auto_parallel_amp.py @@ -877,6 +877,7 @@ class AMPPass(PassBase): naive_set_dist_op_attr_for_program_by_mesh_and_mapping( first_backward_op, ref_mesh, [-1], self.dist_context ) + self._scaled_loss_grad.op = first_backward_op # FIXME(JZ-LIANG) a trick to insert backward op main_block._sync_with_cpp() elementwise_mul_grad_op_desc = main_block.desc._insert_op( diff --git a/python/paddle/distributed/passes/auto_parallel_quantization.py b/python/paddle/distributed/passes/auto_parallel_quantization.py index d63456ab45c2b1f96c8a72a0efee29985d831737..5ece72458abf7fe61c26206f29fbf3cfab84e21f 100644 --- a/python/paddle/distributed/passes/auto_parallel_quantization.py +++ b/python/paddle/distributed/passes/auto_parallel_quantization.py @@ -12,13 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging + +import numpy as np + import paddle -from paddle.distributed.auto_parallel.dist_attribute import ( - OperatorDistributedAttribute, - TensorDistributedAttribute, -) from paddle.fluid import core, framework from paddle.fluid.contrib.slim.quantization import ( + AddQuantDequantForInferencePass, AddQuantDequantPassV2, OutScaleForTrainingPass, QuantizationTransformPassV2, @@ -26,6 +27,11 @@ from paddle.fluid.contrib.slim.quantization import ( ) from paddle.fluid.dygraph.parallel import ParallelEnv +from ..auto_parallel.converter import Converter +from ..auto_parallel.dist_attribute import ( + OperatorDistributedAttribute, + TensorDistributedAttribute, +) from .pass_base import PassBase, register_pass TRANSFORM_PASS_OP_TYPES = utils._weight_supported_quantizable_op_type @@ -42,6 +48,8 @@ class QuantizationPass(PassBase): super().__init__() self.set_attr("dist_context", None) self.set_attr("params_grads", None) + self.set_attr("mode", "train") + self.set_attr("loss", None) def _check_self(self): if self.get_attr("dist_context") is None: @@ -57,15 +65,23 @@ class QuantizationPass(PassBase): dist_context = self.get_attr("dist_context") params_grads = self.get_attr("params_grads") + mode = self.get_attr("mode") + loss = self.get_attr("loss") # TODO: scope and place will be removed, # cause params should be initialized by engine module. scope = paddle.static.global_scope() place = paddle.fluid.CUDAPlace(ParallelEnv().dev_id) + # 0. record the relation among blocks + parent_idx_dict = dict() + for block in main_program.blocks: + parent_idx_dict[block.idx] = block.parent_idx + + is_test = True if mode != "train" else False # 1. Program convert to Graph, and this pass is only for train mode main_graph = framework.IrGraph( - core.Graph(main_program.desc), for_test=False + core.Graph(main_program.desc), for_test=mode != "train" ) # 2. Prepare inputs @@ -91,44 +107,67 @@ class QuantizationPass(PassBase): ) # 3. Add quant op for ops which have parameters - transform_pass = QuantizationTransformPassV2( - scope=scope, - place=place, - weight_bits=self.get_attr('weight_bits'), - activation_bits=self.get_attr('activation_bits'), - skip_pattern=self.get_attr('not_quant_pattern'), - activation_quantize_type="moving_average_abs_max", - quantizable_op_type=transform_pass_ops, - weight_quantize_type=weight_quantize_type, - weight_quantize_func=None, - act_quantize_func=None, - weight_preprocess_func=None, - act_preprocess_func=None, - optimizer_func=None, - executor=None, - ) - transform_pass.apply(main_graph) + if len(transform_pass_ops) > 0: + transform_pass = QuantizationTransformPassV2( + scope=scope, + place=place, + weight_bits=self.get_attr('weight_bits'), + activation_bits=self.get_attr('activation_bits'), + skip_pattern=self.get_attr('not_quant_pattern'), + activation_quantize_type="moving_average_abs_max", + quantizable_op_type=transform_pass_ops, + weight_quantize_type=weight_quantize_type, + weight_quantize_func=None, + act_quantize_func=None, + weight_preprocess_func=None, + act_preprocess_func=None, + optimizer_func=None, + executor=None, + is_test=is_test, + ) + for sub_graph in main_graph.all_sub_graphs(): + transform_pass.apply(sub_graph) # 4. Add quant op for ops which don't have parameter - quant_dequant_pass = AddQuantDequantPassV2( - scope=scope, - place=place, - quant_bits=self.get_attr('activation_bits'), - skip_pattern=self.get_attr('not_quant_pattern'), - quantizable_op_type=quant_dequant_ops, - ) - quant_dequant_pass.apply(main_graph) + if len(quant_dequant_ops) > 0: + quant_dequant_pass = AddQuantDequantPassV2( + scope=scope, + place=place, + quant_bits=self.get_attr('activation_bits'), + skip_pattern=self.get_attr('not_quant_pattern'), + quantizable_op_type=quant_dequant_ops, + is_test=is_test, + ) + for sub_graph in main_graph.all_sub_graphs(): + quant_dequant_pass.apply(sub_graph) # 5. Gather quantitative information for the output out_scale_training_pass = OutScaleForTrainingPass( - scope=scope, place=place + scope=scope, place=place, is_test=is_test ) - out_scale_training_pass.apply(main_graph) + for sub_graph in main_graph.all_sub_graphs(): + out_scale_training_pass.apply(sub_graph) + + # 6. When export quant model, traverse to find the output of each op, and insert the quant/dequant op after it. + if mode != "train" and self.get_attr('onnx_format'): + try: + out_scale_infer_pass = AddQuantDequantForInferencePass( + scope=scope, + place=place, + quant_bits=self.get_attr('activation_bits'), + ) + # for sub_graph in main_graph.all_sub_graphs(): + # out_scale_infer_pass.apply(sub_graph) + except: + logging.warning( + "Unable to convert quant model with onnx_format=True, please update PaddlePaddle >= 2.4.0" + ) - # 6. Convert Graph back to Program + # 7. Convert Graph back to Program quant_program = main_graph.to_program() + quant_program = self.move_presist_var_to_global_block(quant_program) - # 7. get new prams_grads from quant_program + # 8.1 get new prams_grads from quant_program new_params_grads = [] for param, grad in params_grads: if param.name not in quant_program.global_block().vars: @@ -138,7 +177,72 @@ class QuantizationPass(PassBase): new_grad = quant_program.global_block().vars[grad.name] new_params_grads.append((new_param, new_grad)) - # 8. complete distributed attribution + # 8.2 get new loss var + new_loss = None + if loss: + new_loss = quant_program.global_block().vars[loss.name] + + # 8.3 recover the relation among blocks + for block in quant_program.blocks: + block.desc._set_forward_block_idx(parent_idx_dict[block.idx]) + + # 9. complete distributed attribution + self.set_dist_attr_for_qat_program( + quant_program, main_program, dist_context + ) + + # 10. reset scale var value with dist_attr + self.reset_scope_var(quant_program, dist_context, scope, place) + + context.set_attr("main_program", quant_program) + context.set_attr("startup_program", startup_program) + context.set_attr("params_grads", new_params_grads) + context.set_attr("loss", new_loss) + + def move_presist_var_to_global_block(self, program): + global_block = program.global_block() + for _op in global_block.ops: + if _op.type == "while": + _block_id = _op.attr("sub_block").id + _block = program.block(_block_id) + persistables = [] + for _name, _var in _block.vars.items(): + if _var.persistable: + global_block._clone_variable(_var) + persistables.append(_name) + for _name in persistables: + _block._remove_var(_name) + persistables.extend(_op.input('X')) + _op.desc.set_input("X", persistables) + return program + + def reset_scope_var(self, quant_program, dist_context, scope, place): + # The var_value, created by qatization_passes, should has same shape with the value after parallel. + for var in quant_program.list_vars(): + scope_var = scope.find_var(var.name) + if not (scope_var and scope_var.get_tensor()._is_initialized()): + continue + tensor = scope_var.get_tensor() + if var.shape == tensor.shape: + continue + + var_dist_attr = dist_context.get_tensor_dist_attr_for_program(var) + dist_attr = { + "dims_mapping": var_dist_attr.dims_mapping, + "process_shape": var_dist_attr.process_mesh.topology, + "process_group": var_dist_attr.process_mesh.processes, + } + + # slice tensor_value with dist_attr + sliced_tensor = Converter.slice_with_dist_attr( + np.array(tensor), dist_attr + ) + tensor._clear() + tensor.set(sliced_tensor, place) + + def set_dist_attr_for_qat_program( + self, quant_program, main_program, dist_context + ): # NOTE: hack implement, upgrading soon for ib, block in enumerate(quant_program.blocks): # recover origin ops' dist_attr and set quant ops' dist_attr @@ -150,15 +254,22 @@ class QuantizationPass(PassBase): "quantize" in quant_op.type or quant_op.type == "moving_average_abs_max_scale" ): - + # set all quantization ops' dist_attr by quantified op input_name = quant_op.desc.input('X')[0] if "quantize" in input_name: input_name = input_name[ : input_name.index(".quantized") ] - if quant_op.type == "moving_average_abs_max_scale": - consume_op = main_program.blocks[ib].vars[input_name].op + if ( + quant_op.type == "moving_average_abs_max_scale" + or ip - qat_offset >= len(main_program.blocks[ib].ops) + ): + consume_op = ( + main_program.blocks[ib] + ._var_recursive(input_name) + .op + ) else: consume_op = main_program.blocks[ib].ops[ ip - qat_offset @@ -185,23 +296,42 @@ class QuantizationPass(PassBase): ) for slot_name in quant_op.desc.input_names(): + in_name = quant_op.desc.input(slot_name)[0] + input_var = block._var_recursive(in_name) + ref_dims_mapping = [-1] if slot_name == "X": continue - for in_name in quant_op.desc.input(slot_name): - input_var = block.vars[in_name] - tensor_dist_attr = TensorDistributedAttribute() - tensor_dist_attr.process_mesh = ref_process_mesh - tensor_dist_attr.dims_mapping = [-1] - dist_context.set_tensor_dist_attr_for_program( - input_var, tensor_dist_attr - ) - quant_op_dist_attr.set_input_dist_attr( - in_name, tensor_dist_attr - ) + elif slot_name in ['Scale', 'ZeroPoint']: + if ( + quant_op.has_attr('quant_axis') + and quant_op.attr('quant_axis') != -1 + ): + x_name = quant_op.desc.input('X')[0] + x_var = block._var_recursive(x_name) + x_dist_attr = ( + quant_op_dist_attr.get_input_dist_attr( + x_name + ) + ) + quant_axis = quant_op.attr('quant_axis') + ref_dims_mapping = [ + x_dist_attr.dims_mapping[quant_axis] + ] + + tensor_dist_attr = TensorDistributedAttribute() + tensor_dist_attr.process_mesh = ref_process_mesh + tensor_dist_attr.dims_mapping = ref_dims_mapping + dist_context.set_tensor_dist_attr_for_program( + input_var, tensor_dist_attr + ) + quant_op_dist_attr.set_input_dist_attr( + in_name, tensor_dist_attr + ) for slot_name in quant_op.desc.output_names(): output_name = quant_op.desc.output(slot_name)[0] - output_var = block.vars[output_name] + output_var = block._var_recursive(output_name) + ref_dims_mapping = [-1] if slot_name == "Y": dist_context.set_tensor_dist_attr_for_program( output_var, consume_input_dist_attr @@ -209,22 +339,39 @@ class QuantizationPass(PassBase): quant_op_dist_attr.set_output_dist_attr( output_name, consume_input_dist_attr ) - else: - tensor_dist_attr = TensorDistributedAttribute() - tensor_dist_attr.process_mesh = ref_process_mesh - tensor_dist_attr.dims_mapping = [-1] - dist_context.set_tensor_dist_attr_for_program( - output_var, tensor_dist_attr - ) - quant_op_dist_attr.set_output_dist_attr( - output_name, tensor_dist_attr - ) + continue + elif slot_name == "OutScale": + if ( + quant_op.has_attr('quant_axis') + and quant_op.attr('quant_axis') != -1 + ): + x_name = quant_op.desc.input('X')[0] + x_var = block._var_recursive(x_name) + x_dist_attr = ( + quant_op_dist_attr.get_input_dist_attr( + x_name + ) + ) + quant_axis = quant_op.attr('quant_axis') + ref_dims_mapping = [ + x_dist_attr.dims_mapping[quant_axis] + ] + + tensor_dist_attr = TensorDistributedAttribute() + tensor_dist_attr.process_mesh = ref_process_mesh + tensor_dist_attr.dims_mapping = ref_dims_mapping + dist_context.set_tensor_dist_attr_for_program( + output_var, tensor_dist_attr + ) + quant_op_dist_attr.set_output_dist_attr( + output_name, tensor_dist_attr + ) quant_op._set_attr("op_device", "") qat_offset += 1 else: - + # recover origin ops' dist_attr origin_op = main_program.blocks[ib].ops[ip - qat_offset] quant_op.desc.set_original_id(origin_op.desc.original_id()) dist_origin_op = dist_context.get_dist_op_for_program( @@ -240,7 +387,21 @@ class QuantizationPass(PassBase): quant_op_dist_attr.process_mesh = ( origin_op_dist_attr.process_mesh ) + + scale_offset = 0 for idx, input_name in enumerate(quant_op.input_arg_names): + if ( + origin_op.type == "while" + and input_name not in origin_op.input_arg_names + ): + assert ( + "@scale" in input_name + or "@zero_point" in input_name + ) + scale_offset += 1 + continue + + idx -= scale_offset origin_input_name = origin_op.input_arg_names[idx] origin_input_dist_attr = ( origin_op_dist_attr.inputs_dist_attrs[ @@ -251,20 +412,6 @@ class QuantizationPass(PassBase): input_name, origin_input_dist_attr ) - if input_name not in main_program.blocks[ib].vars: - origin_input_var = main_program.blocks[ib].vars[ - origin_input_name - ] - origin_in_tensor_dist_attr = ( - dist_context.get_dist_tensor_for_program( - origin_input_var - ).dist_attr - ) - quant_input_var = block.vars[input_name] - dist_context.set_tensor_dist_attr_for_program( - quant_input_var, origin_in_tensor_dist_attr - ) - for idx, output_name in enumerate( quant_op.output_arg_names ): @@ -278,16 +425,18 @@ class QuantizationPass(PassBase): output_name, origin_output_dist_attr ) - if output_name not in main_program.blocks[ib].vars: - origin_output_var = main_program.blocks[ib].vars[ - origin_output_name - ] + if not main_program.blocks[ib]._find_var_recursive( + output_name + ): + origin_output_var = main_program.blocks[ + ib + ]._var_recursive(origin_output_name) origin_out_tensor_dist_attr = ( dist_context.get_dist_tensor_for_program( origin_output_var ).dist_attr ) - quant_output_var = block.vars[output_name] + quant_output_var = block._var_recursive(output_name) dist_context.set_tensor_dist_attr_for_program( quant_output_var, origin_out_tensor_dist_attr ) @@ -308,7 +457,3 @@ class QuantizationPass(PassBase): dist_context.set_tensor_dist_attr_for_program( dst_var, dist_tensor.dist_attr ) - - context.set_attr("main_program", quant_program) - context.set_attr("startup_program", startup_program) - context.set_attr("params_grads", new_params_grads) diff --git a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py index 0a14b6326674ba65a6175e20d66e9f8de557bf3d..0b1e3cd96d7d604310dc35496fc1214c9243add9 100644 --- a/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py +++ b/python/paddle/fluid/contrib/slim/quantization/quantization_pass.py @@ -513,11 +513,12 @@ class QuantizationTransformPass: var_dtype=var_node.dtype(), ) scale_name = self._quantized_scale_name(name) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" try: scale_value = np.array( self._scope.find_var(scale_name).get_tensor() @@ -560,11 +561,12 @@ class QuantizationTransformPass: ) scale_name = self._quantized_scale_name(name) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" try: scale_value = np.array( self._scope.find_var(scale_name).get_tensor() @@ -591,11 +593,12 @@ class QuantizationTransformPass: shape=[self._window_size], var_dtype=var_node.dtype(), ) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" _init_var_node( scales_node, np.zeros([self._window_size], dtype=data_type), @@ -640,11 +643,12 @@ class QuantizationTransformPass: var_dtype=var_node.dtype(), ) scale_name = self._quantized_scale_name(name) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" try: scale_value = np.array( self._scope.find_var(scale_name).get_tensor() @@ -669,11 +673,12 @@ class QuantizationTransformPass: var_dtype=var_node.dtype(), shape=[1], ) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" _init_var_node( state_in_node, np.ones([1], dtype=data_type), @@ -746,11 +751,12 @@ class QuantizationTransformPass: var_dtype=var_node.dtype(), ) scale_name = self._quantized_scale_name(name) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" try: scale_value = np.array( self._scope.find_var(scale_name).get_tensor() @@ -1287,11 +1293,13 @@ class QuantizationFreezePass: shape=[channel_scale.shape[0]], var_dtype=output_var_node.dtype(), ) - data_type = ( - 'float64' - if output_var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + + if output_var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif output_var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" _init_var_node( weight_scale_node, channel_scale.astype(data_type), @@ -1439,6 +1447,7 @@ class QuantizationFreezePass: def _is_float(self, v): return ( isinstance(v, float) + or isinstance(v, np.float16) or isinstance(v, np.float32) or isinstance(v, np.float64) ) @@ -1636,14 +1645,17 @@ class OutScaleForTrainingPass: if in_node.dtype() not in [ core.VarDesc.VarType.FP64, core.VarDesc.VarType.FP32, + core.VarDesc.VarType.FP16, ]: continue - data_type = ( - 'float64' - if in_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if in_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif in_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" + try: graph._find_node_by_name( graph.all_var_nodes(), @@ -1781,6 +1793,7 @@ class OutScaleForInferencePass: not in [ core.VarDesc.VarType.FP64, core.VarDesc.VarType.FP32, + core.VarDesc.VarType.FP16, ] ): continue @@ -1997,11 +2010,12 @@ class AddQuantDequantPass: var_dtype=var_node.dtype(), ) scale_name = "{}.quant_dequant@scale".format(var_node.name()) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" try: if ( self._scale_dict is not None @@ -2036,11 +2050,12 @@ class AddQuantDequantPass: var_dtype=var_node.dtype(), shape=[1], ) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" _init_var_node( state_in_node, np.ones([1], dtype=data_type), @@ -2149,11 +2164,12 @@ class InsertQuantizeLinear: var_dtype=var_node.dtype(), ) if not scale_var_node: - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" scale_name = self._quantized_scale_name(var_name) if self.channel_wise: scale_var_shape = var_node.shape()[self.quant_axis] @@ -2218,11 +2234,12 @@ class InsertQuantizeLinear: var_dtype=var_node.dtype(), shape=[1], ) - data_type = ( - 'float64' - if var_node.dtype() == core.VarDesc.VarType.FP64 - else 'float32' - ) + if var_node.dtype() == core.VarDesc.VarType.FP64: + data_type = 'float64' + elif var_node.dtype() == core.VarDesc.VarType.FP32: + data_type = 'float32' + else: + data_type = "float16" _init_var_node( state_in_node, np.ones([1], dtype=data_type), @@ -3277,6 +3294,7 @@ class AddQuantDequantForInferencePass: if out_node.dtype() not in [ core.VarDesc.VarType.FP64, core.VarDesc.VarType.FP32, + core.VarDesc.VarType.FP16, ]: continue if var_name in dequantized_vars_map: diff --git a/python/paddle/fluid/tests/unittests/auto_parallel/CMakeLists.txt b/python/paddle/fluid/tests/unittests/auto_parallel/CMakeLists.txt index 21c0f88438ad818a17e48195438feb30bb3b538c..d3570920a5842cc9a3167550613c1bcf633219d8 100644 --- a/python/paddle/fluid/tests/unittests/auto_parallel/CMakeLists.txt +++ b/python/paddle/fluid/tests/unittests/auto_parallel/CMakeLists.txt @@ -74,6 +74,9 @@ if(WITH_DISTRIBUTE AND WITH_GPU) set_tests_properties(test_parallel_tuner_predict PROPERTIES TIMEOUT 120) py_test_modules(test_selective_recompute MODULES test_selective_recompute) set_tests_properties(test_selective_recompute PROPERTIES TIMEOUT 50) + py_test_modules(test_pass_quantization MODULES test_pass_quantization ENVS + ${dist_ENVS}) + set_tests_properties(test_pass_quantization PROPERTIES TIMEOUT 60) py_test_modules(test_tuning_recompute MODULES test_tuning_recompute) set_tests_properties(test_tuning_recompute PROPERTIES TIMEOUT 240) @@ -113,7 +116,6 @@ if(WITH_DISTRIBUTE AND WITH_GPU) py_test_modules(test_process_mesh MODULES test_process_mesh) py_test_modules(test_interface MODULES test_interface) py_test_modules(test_strategy MODULES test_strategy) - py_test_modules(test_pass_quantization MODULES test_pass_quantization) py_test_modules(test_dist_shape MODULES test_dist_shape) py_test_modules(test_dist_assign MODULES test_dist_assign) py_test_modules(test_conditional_block_reshard MODULES diff --git a/python/paddle/fluid/tests/unittests/auto_parallel/quantization_pass_unittest.py b/python/paddle/fluid/tests/unittests/auto_parallel/quantization_pass_unittest.py new file mode 100644 index 0000000000000000000000000000000000000000..8b941e7a7fc6cc68e48d24828b81a462737bea10 --- /dev/null +++ b/python/paddle/fluid/tests/unittests/auto_parallel/quantization_pass_unittest.py @@ -0,0 +1,167 @@ +# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile +import unittest + +import numpy as np +from get_gpt_model import FakeDataset, create_data_holder, generate_model + +import paddle +from paddle.distributed.fleet import auto + +paddle.enable_static() + + +def apply_pass(): + dist_strategy = auto.Strategy() + dist_strategy.auto_mode = "semi" + + amp = dist_strategy.amp + amp.enable = True + amp.custom_white_list = ["lookup_table", "lookup_table_v2"] + amp.custom_black_list = [ + "reduce_sum", + "c_softmax_with_cross_entropy", + "elementwise_div", + ] + amp.init_loss_scaling = 32768 + amp.use_fp16_guard = False + amp.use_pure_fp16 = True + + qat = dist_strategy.qat + qat.enable = True + qat.channel_wise_abs_max = True + qat.weight_bits = 8 + qat.activation_bits = 8 + qat.not_quant_pattern = ['skip_quant'] + qat.onnx_format = True + return dist_strategy + + +class TestQuantizationPassTrain(unittest.TestCase): + def test_qat_pass_training(self): + + batch_size = 1 + batch_num = 10 + + strategy = apply_pass() + model, loss = generate_model("mp") + opt = paddle.optimizer.AdamW(learning_rate=0.00001) + engine = auto.Engine(model, loss, opt, strategy=strategy) + dataset = FakeDataset(batch_size * batch_num) + engine.fit(dataset, 3, batch_size=batch_size) + self.check_program(engine.main_program) + + def check_program(self, program): + + quantizable_op_and_inputs = {'matmul_v2': ['X', 'Y']} + quantizable_grad_op_inputs = {'matmul_v2_grad': ['X', 'Y']} + + quantized_ops = set() + for block in program.blocks: + for idx, op in enumerate(block.ops): + is_quntized = False + if op.type in quantizable_op_and_inputs: + for arg_name in op.input_arg_names: + if ".quantized" in arg_name: + is_quntized = True + + if not is_quntized: + continue + + # check forward + if op.type in quantizable_op_and_inputs: + + for arg_name in op.input_arg_names: + if "c_identity" in arg_name: + arg_name = block.ops[idx - 1].input_arg_names[0] + assert arg_name.endswith('.quantized.dequantized') + quantized_ops.add(arg_name) + + for op in block.ops: + is_quntized = False + if op.type in quantizable_grad_op_inputs: + for pname in quantizable_grad_op_inputs[op.type]: + arg_name = op.input(pname)[0] + if ".quantized" in arg_name: + is_quntized = True + + if not is_quntized: + continue + + # check backward + if op.type in quantizable_grad_op_inputs: + for pname in quantizable_grad_op_inputs[op.type]: + arg_name = op.input(pname)[0] + assert arg_name.endswith('.quantized.dequantized') + assert arg_name in quantized_ops + + +class TestQuantizationPassExport(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_qat_pass_2(self): + + batch_size = 1 + batch_num = 10 + + strategy = apply_pass() + model, loss = generate_model("mp") + engine = auto.Engine(model, loss, strategy=strategy) + inputs_spec, labels_spec = create_data_holder(batch_size=1) + engine.prepare(inputs_spec, labels_spec, mode="predict") + + path = os.path.join(self.temp_dir.name, 'inf') + engine.save(path, training=False) + self.check_export(engine._executor) + + def check_export(self, exe): + sequence_len = 512 + vocab_size = 1000 + + tokens = [np.random.randint(vocab_size, size=sequence_len)] + position_ids = [np.arange(sequence_len)] + attention_mask = [np.tril(np.ones(sequence_len))] + + path_prefix = os.path.join( + self.temp_dir.name, + 'inf_dist{}'.format(paddle.distributed.get_rank()), + ) + [ + inference_program, + feed_target_names, + fetch_targets, + ] = paddle.static.load_inference_model( + path_prefix=path_prefix, executor=exe + ) + + out = exe.run( + inference_program, + feed={ + "tokens": tokens, + "position_ids": position_ids, + "attention_mask": attention_mask, + }, + fetch_list=fetch_targets, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/paddle/fluid/tests/unittests/auto_parallel/test_pass_quantization.py b/python/paddle/fluid/tests/unittests/auto_parallel/test_pass_quantization.py index e87af55f63cc7da5568402c177e44566f1dd7157..4e4a1fa724d7cd32a1eb15b7ee17adcb2c23840a 100644 --- a/python/paddle/fluid/tests/unittests/auto_parallel/test_pass_quantization.py +++ b/python/paddle/fluid/tests/unittests/auto_parallel/test_pass_quantization.py @@ -12,83 +12,45 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import subprocess +import sys +import tempfile import unittest -from get_gpt_model import FakeDataset, generate_model - -import paddle -from paddle.distributed.fleet import auto - -paddle.enable_static() - - -def apply_pass(): - dist_strategy = auto.Strategy() - dist_strategy.auto_mode = "semi" - qat = dist_strategy.qat - qat.enable = True - qat.channel_wise_abs_max = True - qat.weight_bits = 8 - qat.activation_bits = 8 - qat.not_quant_pattern = ['skip_quant'] - return dist_strategy - class TestQuantizationPass(unittest.TestCase): - def test_qat_pass(self): - - batch_size = 8 - batch_num = 10 - - strategy = apply_pass() - model, loss = generate_model("serial") - opt = paddle.optimizer.AdamW(learning_rate=0.00001) - engine = auto.Engine(model, loss, opt, strategy=strategy) - dataset = FakeDataset(batch_size * batch_num) - engine.fit(dataset, 3, batch_size=batch_size) - - self.check_program(engine.main_program) - - def check_program(self, program): - - quantizable_op_and_inputs = {'matmul_v2': ['X', 'Y']} - quantizable_grad_op_inputs = {'matmul_v2_grad': ['X', 'Y']} - - quantized_ops = set() - for block in program.blocks: - for op in block.ops: - is_quntized = False - if op.type in quantizable_op_and_inputs: - for arg_name in op.input_arg_names: - if ".quantized" in arg_name: - is_quntized = True - - if not is_quntized: - continue - - # check forward - if op.type in quantizable_op_and_inputs: - for arg_name in op.input_arg_names: - assert arg_name.endswith('.quantized.dequantized') - quantized_ops.add(arg_name) - - for op in block.ops: - is_quntized = False - if op.type in quantizable_grad_op_inputs: - for pname in quantizable_grad_op_inputs[op.type]: - arg_name = op.input(pname)[0] - if ".quantized" in arg_name: - is_quntized = True - - if not is_quntized: - continue - - # check backward - if op.type in quantizable_grad_op_inputs: - for pname in quantizable_grad_op_inputs[op.type]: - arg_name = op.input(pname)[0] - assert arg_name.endswith('.quantized.dequantized') - assert arg_name in quantized_ops + def test_mp2(self): + file_dir = os.path.dirname(os.path.abspath(__file__)) + launch_model_path = os.path.join( + file_dir, "quantization_pass_unittest.py" + ) + + if os.environ.get("WITH_COVERAGE", "OFF") == "ON": + coverage_args = ["-m", "coverage", "run", "--branch", "-p"] + else: + coverage_args = [] + + tmp_dir = tempfile.TemporaryDirectory() + cmd = ( + [sys.executable, "-u"] + + coverage_args + + [ + "-m", + "paddle.distributed.launch", + "--devices", + "0,1", + "--log_dir", + tmp_dir.name, + launch_model_path, + ] + ) + + process = subprocess.Popen(cmd) + process.wait() + self.assertEqual(process.returncode, 0) + + tmp_dir.cleanup() if __name__ == "__main__": diff --git a/python/paddle/fluid/tests/unittests/auto_parallel_gpt_model.py b/python/paddle/fluid/tests/unittests/auto_parallel_gpt_model.py index 2edb5360e471d6a1805346513f7b7279bcc7166f..f0e55ba21d59bc04f0a4dbfd65ce2d4cec08b87d 100644 --- a/python/paddle/fluid/tests/unittests/auto_parallel_gpt_model.py +++ b/python/paddle/fluid/tests/unittests/auto_parallel_gpt_model.py @@ -834,13 +834,14 @@ class GPTForPretraining(nn.Layer): x_dims_mapping = ["x"] + [None for i in range(len(x.shape) - 1)] w_dims_mapping = ["y"] + [None for i in range(len(w.shape) - 1)] - if mesh: - matmul = auto.shard_op( - paddle.matmul, mesh, [x_dims_mapping, w_dims_mapping, None] - ) - logits = matmul(x, w, transpose_y=True) - else: - logits = paddle.matmul(x, w, transpose_y=True) + with paddle.fluid.name_scope('skip_quant'): + if mesh: + matmul = auto.shard_op( + paddle.matmul, mesh, [x_dims_mapping, w_dims_mapping, None] + ) + logits = matmul(x, w, transpose_y=True) + else: + logits = paddle.matmul(x, w, transpose_y=True) if use_cache: return logits, cached_kvs diff --git a/python/paddle/fluid/tests/unittests/test_fake_dequantize_op.py b/python/paddle/fluid/tests/unittests/test_fake_dequantize_op.py index 47ab4754f2a3f07f82158a542e11fc5d1e296415..f9824664930fba02fc3b59cab3cc4d9b048f12ab 100644 --- a/python/paddle/fluid/tests/unittests/test_fake_dequantize_op.py +++ b/python/paddle/fluid/tests/unittests/test_fake_dequantize_op.py @@ -26,7 +26,7 @@ def quantize_max_abs(x, max_range): def dequantize_max_abs(x, scale, max_range): - y = (scale / max_range) * x + y = x * scale / max_range return y @@ -292,6 +292,45 @@ class TestDequantizeOpDouble(TestDequantizeOp): self.quant_axis = -1 +class TestDequantizeOpHalf(TestDequantizeOp): + def set_args(self): + self.bit_length = 8 + self.max_range = math.pow(2, self.bit_length - 1) - 1 + self.data_type = "float16" + self.quant_axis = -1 + + def setUp(self): + self.set_args() + self.op_type = "dequantize_linear" + x = np.random.randn(31, 65).astype(np.float16) + yq, scale = quantize_max_abs(x, self.max_range) + scale = np.array(scale).astype('float16') + yq = np.array(yq).astype('int8') + ydq = dequantize_max_abs(yq, scale, self.max_range) + ydq = ydq.astype('float16') + zero_point = np.zeros(scale.shape, dtype="int32") + + self.inputs = {'X': yq, 'Scale': scale, 'ZeroPoint': zero_point} + self.attrs = { + 'bit_length': self.bit_length, + 'quant_axis': self.quant_axis, + } + self.outputs = {'Y': ydq} + + def _get_places(self): + import paddle + import paddle.fluid.core as core + + if core.is_compiled_with_cuda(): + place = paddle.fluid.core.CUDAPlace(0) + if paddle.fluid.core.is_float16_supported(place): + return [place] + else: + return [] + else: + return [] + + class TestDequantizeOp5Bits(TestDequantizeOp): def set_args(self): self.bit_length = 5