# Copyright (c) 2020 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 import logging from functools import reduce from .meta_optimizer_base import MetaOptimizerBase __all__ = [] import paddle from paddle.base import framework from paddle.base.dygraph import base as imperative_base from paddle.common_ops_import import LayerHelper from paddle.framework import core, in_dynamic_mode from paddle.nn.clip import ClipGradByNorm, append_gradient_clip_ops from paddle.optimizer import Momentum, Optimizer from paddle.regularizer import L1Decay, L2Decay from paddle.static import create_global_var class DGCMomentumOptimizer(Optimizer): _u_velocity_acc_str = "_dgc_u_" _v_velocity_acc_str = "_dgc_v_" def __init__( self, learning_rate, momentum, rampup_begin_step, rampup_step=1, sparsity=[0.999], parameter_list=None, use_nesterov=False, num_trainers=None, regularization=None, grad_clip=None, name=None, ): if in_dynamic_mode(): raise Exception("In dygraph, don't support DGCMomentumOptimizer.") assert core.is_compiled_with_cuda(), ( "Paddle is not compiled with CUDA. DGC is only support GPU for now." ) assert learning_rate is not None assert momentum is not None super().__init__( learning_rate=learning_rate, parameters=parameter_list, weight_decay=regularization, grad_clip=grad_clip, name=name, ) self.type = "dgc_momentum" self._momentum = momentum self._use_nesterov = bool(use_nesterov) assert rampup_begin_step >= 0, "rampup_begin_step must >= 0" self._rampup_begin_step = rampup_begin_step self._rampup_step = rampup_step self._sparsity = sparsity self._rampup_begin_step_var = None self._global_step_var = None self._dgc_clip_norm = None self._num_trainers = num_trainers if grad_clip is not None: if not isinstance(grad_clip, ClipGradByNorm): raise TypeError( "The type of grad_clip should be 'ClipGradByNorm', because DGCMomentumOptimizer only support ClipGradByNorm" ) assert isinstance(num_trainers, int), ( f"The type of num_trainers should be 'int', but received {type(num_trainers)}" ) assert num_trainers > 0, ( "The value of num_trainers should be greater than 0!" ) self._dgc_clip_norm = grad_clip.clip_norm * (num_trainers**-0.5) self.regular_type, self.regular_coeff = self._get_regularization_param( self.regularization ) def _get_regularization_param(self, regularization): regular_type = 0 regular_coeff = 0.0 if regularization is not None: regular_coeff = regularization._coeff if isinstance(regularization, L1Decay): regular_type = 1 elif isinstance(regularization, L2Decay): regular_type = 2 else: raise AssertionError( "regularization must be None|L1Decay|L2Deacy" ) return regular_type, regular_coeff def _is_use_dgc(self, param_var, grad_var): var_numel = abs(reduce(lambda x, y: x * y, param_var.shape, 1)) if ( var_numel < 16384 or param_var.type == core.VarDesc.VarType.SELECTED_ROWS or grad_var.type == core.VarDesc.VarType.SELECTED_ROWS or param_var.dtype != core.VarDesc.VarType.FP32 ): return False return True def _append_optimize_op(self, block, param_and_grad): assert isinstance(block, paddle.framework.Block) velocity_acc = self._get_accumulator( self._u_velocity_acc_str, param_and_grad[0] ) assert velocity_acc is not None inputs = { "Param": param_and_grad[0], "Grad": param_and_grad[1], "Velocity": velocity_acc, "LearningRate": self._create_param_lr(param_and_grad), } outputs = { "ParamOut": param_and_grad[0], "VelocityOut": velocity_acc, } attrs = {"mu": self._momentum, "use_nesterov": self._use_nesterov} if not self._is_use_dgc(param_and_grad[0], param_and_grad[1]): type = "momentum" else: type = "dgc_momentum" inputs.update( { "current_step": self._global_step_var, "nranks": self._nranks_var, } ) outputs.update({'Grad_out': param_and_grad[1]}) attrs.update({"rampup_begin_step": float(self._rampup_begin_step)}) # create the dgc momentum optimize op dgc_momentum_op = block.append_op( type=type, inputs=inputs, outputs=outputs, attrs=attrs, stop_gradient=True, ) return dgc_momentum_op def _add_auto_increment_var(self, counter_name, begin, step=1): helper = LayerHelper('global_step_counter') counter, is_new_var = helper.create_or_get_global_variable( name=counter_name, dtype='float32', shape=[1], persistable=True ) if is_new_var: helper.set_variable_initializer( counter, initializer=paddle.nn.initializer.ConstantInitializer( value=float(begin - 1), force_cpu=True ), ) helper.main_program.global_block()._prepend_op( type='increment', inputs={'X': [counter]}, outputs={'Out': [counter]}, attrs={'step': float(step)}, stop_gradient=True, ) counter.stop_gradient = True return counter def _add_nranks_var(self, name, value=-1): helper = LayerHelper('global_step_counter') counter, is_new_var = helper.create_or_get_global_variable( name=name, dtype='float32', shape=[1], persistable=True ) if is_new_var: helper.set_variable_initializer( counter, initializer=paddle.nn.initializer.ConstantInitializer( value=float(value), force_cpu=True ), ) counter.stop_gradient = True return counter def _append_dgc_ops(self, param_and_grads): main_program = paddle.static.default_main_program() main_program._enable_dgc = True # step counter self._global_step_var = self._add_auto_increment_var( counter_name=core.dgc.kDGCCounterName(), begin=0 ) self._nranks_var = self._add_nranks_var( name=core.dgc.kDGCNRanksName(), value=self._num_trainers ) # rampup begin step var for all_reduce_op_handle self._rampup_begin_step_var = create_global_var( shape=[1], dtype=core.VarDesc.VarType.FP32, persistable=True, name=core.dgc.kDGCRampUpBeginStepName(), value=self._rampup_begin_step * 1.0, force_cpu=True, ) self.helper = LayerHelper(self.__class__.__name__) for param_var, grad_var in param_and_grads: # reuse velocity in dgc_op and dgc_momentum_op u_var = self._add_accumulator(self._u_velocity_acc_str, param_var) if not self._is_use_dgc(param_var, grad_var): continue v_var = self._add_accumulator(self._v_velocity_acc_str, param_var) k_var = create_global_var( shape=[1], dtype=param_var.dtype, persistable=True, name=param_var.name + core.dgc.kDGCKName(), value=0.0, force_cpu=True, ) encoded_var = create_global_var( shape=[1], dtype=param_var.dtype, persistable=True, name=param_var.name + core.dgc.kDGCEncodedName(), value=0.0, force_cpu=False, ) gather_var = create_global_var( shape=[1], dtype=param_var.dtype, persistable=True, name=param_var.name + core.dgc.kDGCGatherName(), value=0.0, force_cpu=False, ) # del back oprolevarname op_maker = core.op_proto_and_checker_maker backward = core.op_proto_and_checker_maker.OpRole.Backward for op in main_program.global_block().ops: if not self._is_the_backward_op(op): continue var_attr = op.all_attrs()[op_maker.kOpRoleVarAttrName()] if param_var.name not in var_attr: continue var_attr.remove(param_var.name) var_attr.remove(grad_var.name) if len(var_attr) > 1: op._set_attr(op_maker.kOpRoleVarAttrName(), var_attr) else: op._remove_attr(op_maker.kOpRoleVarAttrName()) clip_var = grad_var if self._dgc_clip_norm is not None: clip_var = self._append_clip_norm(grad_var, self._dgc_clip_norm) self._dgc_op( param_var, clip_var, grad_var, u_var, v_var, k_var, encoded_var, gather_var, ) def _is_the_backward_op(self, op): op_maker = core.op_proto_and_checker_maker backward = core.op_proto_and_checker_maker.OpRole.Backward if op_maker.kOpRoleVarAttrName() in op.attr_names and int( op.all_attrs()[op_maker.kOpRoleAttrName()] ) == int(backward): return True return False def _clip_by_norm(self, x, max_norm, name=None): args = {'x': x, 'max_norm': max_norm, 'name': name} helper = LayerHelper("dgc_clip_by_norm_op", **args) if name is None: name = paddle.base.unique_name.generate_with_ignorable_key( ".".join([helper.name, 'tmp']) ) out = helper.create_variable( type=x.type, name=name, dtype=x.dtype, persistable=False ) helper.append_op( type="dgc_clip_by_norm", inputs={"X": x, "current_step": self._global_step_var}, attrs={ "max_norm": max_norm, "rampup_begin_step": float(self._rampup_begin_step), }, outputs={"Out": out}, ) return out def _append_clip_norm(self, grad_var, clip_norm): with grad_var.block.program._backward_role_guard(): return self._clip_by_norm( x=grad_var, max_norm=clip_norm, name=grad_var.name ) def _dgc_op( self, param_var, clip_var, grad_var, u_var, v_var, k_var, encoded_var, gather_var, ): block = paddle.static.default_main_program().global_block() op_maker = core.op_proto_and_checker_maker regular_type = self.regular_type regular_coeff = self.regular_coeff # The regularizer of the Parameters have higher priority if param_var.regularizer is not None: regular_type, regular_coeff = self._get_regularization_param( param_var.regularizer ) dgc_op = block.append_op( type="dgc", inputs={ "U": u_var, "V": v_var, "Grad": clip_var, "Param": param_var, "current_step": self._global_step_var, "nranks": self._nranks_var, }, outputs={ "U_out": u_var, "V_out": v_var, "EncodeGrad": encoded_var, "k": k_var, "Grad_out": grad_var, "GatherBuff": gather_var, }, attrs={ "m": self._momentum, "sparsity": self._sparsity, "use_nesterov": self._use_nesterov, "rampup_begin_step": float(self._rampup_begin_step), "rampup_step": float(self._rampup_step), "regular_coeff": float(regular_coeff), "regular_type": int(regular_type), }, stop_gradient=True, ) backward = op_maker.OpRole.Backward dgc_op._set_attr(op_maker.kOpRoleAttrName(), backward) dgc_op._set_attr( op_maker.kOpRoleVarAttrName(), [param_var.name, grad_var.name] ) def _process_distribute_lookuptable(self, param_grads): """ Because distribute lookup table only support SGD optimizer for now, not support other optimizer and regularization, so we should find the table parameter out, and avoid to add regularization and other op for it, and add sgd optimize op for it independently. :param param_grads(list((Var, Var))): list of (param, grad) pair. :param loss: the loss variable. :param startup_program: the startup program """ from paddle.distributed.distribute_lookup_table import ( find_distributed_lookup_table, ) program = framework.default_main_program() global_block = framework.default_main_program().global_block() table_name = find_distributed_lookup_table(program) table_param = None table_grad = None new_param_grads = [] for p, g in param_grads: if p.name == table_name: if table_param is not None: raise RuntimeError( "multi dist table var found, only support one now!" ) table_param = p table_grad = g else: new_param_grads.append((p, g)) sgd_op = None if table_param is not None: param_and_grad = [table_param, table_grad] with ( table_param.block.program._optimized_guard(param_and_grad), framework.name_scope("optimizer"), ): self._create_global_learning_rate() # create the optimize op sgd_op = global_block.append_op( type='sgd', inputs={ "Param": table_param, "Grad": table_grad, "LearningRate": self._create_param_lr(param_and_grad), }, outputs={"ParamOut": param_and_grad[0]}, ) return new_param_grads, (table_param, table_grad), sgd_op @imperative_base.no_grad() def apply_gradients(self, params_grads): # Note: since we can't use all_reduce_op now, # dgc_op should be the last op of one grad. # Maybe need a grad allreduce pass. self._append_dgc_ops(params_grads) params_grads = sorted(params_grads, key=lambda x: x[0].name) ( params_grads, table_param_and_grad, table_optimize_op, ) = self._process_distribute_lookuptable(params_grads) not_dgc_params_grads = [] dgc_params_grads = [] # DGC clip and regularization in optimizer.backward for param, grad in params_grads: if not self._is_use_dgc(param, grad): not_dgc_params_grads.append((param, grad)) else: dgc_params_grads.append((param, grad)) # 'optimizer(grad_clip)' or 'set_gradient_clip' if self._grad_clip is not None: not_dgc_params_grads = self._grad_clip(not_dgc_params_grads) else: not_dgc_params_grads = append_gradient_clip_ops( not_dgc_params_grads ) not_dgc_params_grads = self.append_regularization_ops( not_dgc_params_grads, self.regularization ) params_grads = not_dgc_params_grads + dgc_params_grads params_grads = sorted(params_grads, key=lambda x: x[0].name) optimize_ops = self._create_optimization_pass(params_grads) if table_optimize_op is not None: optimize_ops.append(table_optimize_op) params_grads.append(table_param_and_grad) return optimize_ops class DGCOptimizer(MetaOptimizerBase): def __init__(self, optimizer): super().__init__(optimizer) self.inner_opt = optimizer self.dgc_opt = None # we do not allow meta optimizer to be inner optimizer currently self.meta_optimizers_white_list = [] self.meta_optimizers_black_list = [] def _set_basic_info( self, loss, role_maker, user_defined_optimizer, user_defined_strategy ): super()._set_basic_info( loss, role_maker, user_defined_optimizer, user_defined_strategy ) def _init_dgc_opt(self): if self.dgc_opt is not None: return opt = self.inner_opt if not self.role_maker._is_collective: return if not isinstance(opt, Momentum): return configs = self.user_defined_strategy.dgc_configs if len(configs['sparsity']) == 0: # default is [0.999] configs['sparsity'] = [0.999] self.dgc_opt = DGCMomentumOptimizer( learning_rate=opt._learning_rate, momentum=opt._momentum, rampup_begin_step=configs['rampup_begin_step'], rampup_step=configs['rampup_step'], sparsity=configs['sparsity'], parameter_list=opt._parameter_list, use_nesterov=opt._use_nesterov, num_trainers=self.role_maker._worker_num(), regularization=opt.regularization, grad_clip=opt._grad_clip, name=opt._name, ) def _can_apply(self): if not self.role_maker._is_collective: return False if self.user_defined_strategy.dgc: if not isinstance(self.inner_opt, Momentum): logging.warning("dgc only works on Momentum optimizer") return False if self.role_maker._worker_num() <= 1: logging.warning("dgc only works on multi cards") return False return True return False def _disable_strategy(self, dist_strategy): dist_strategy.dgc = False dist_strategy.dgc_configs = {} def _enable_strategy(self, dist_strategy, context): dist_strategy.dgc = True dist_strategy.dgc_configs = {"rampup_begin_step": 0, "rampup_step": 1} def backward( self, loss, startup_program=None, parameter_list=None, no_grad_set=None, callbacks=None, ): self._init_dgc_opt() return self.dgc_opt.backward( loss, startup_program, parameter_list, no_grad_set, callbacks ) def apply_gradients(self, params_grads): self._init_dgc_opt() return self.dgc_opt.apply_gradients(params_grads=params_grads) def apply_optimize(self, loss, startup_program, params_grads): self._init_dgc_opt() return self.dgc_opt._apply_optimize( loss, startup_program=startup_program, params_grads=params_grads ) def minimize_impl( self, loss, startup_program=None, parameter_list=None, no_grad_set=None ): self._init_dgc_opt() optimize_ops, params_grads = self.dgc_opt.minimize( loss, startup_program, parameter_list, no_grad_set ) return optimize_ops, params_grads