# 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 copy import inspect import logging from collections import defaultdict import paddle from paddle import core from paddle.jit import not_to_static, to_static from paddle.jit.dy2static.program_translator import ( ProgramTranslator, StaticFunction, ) from paddle.jit.dy2static.utils import as_not_paddle_func from paddle.nn import Layer from paddle.static import Parameter, global_scope, program_guard from paddle.static.amp.fp16_utils import ( DEFAULT_AMP_OPTIONS, prepare_op_amp_options, ) from .converter import Converter from .dist_attribute import TensorDistAttr from .process_group import get_world_process_group from .utils import get_logger, to_list class ProxyLayer(Layer): """ ProxyLayer implements all logic for converting dygraph model into static Program IR. Meanwhile, it provides conventional interfaces for auto parallel to visit feed/fetch/loss/metric variables. """ def __init__(self, layer, loss_func, metrics): super().__init__() # NOTE: All verify logics are finished in Engine.Prepare self.inner_layer = layer self.loss_func = loss_func self.metrics = metrics # train / eval / predict self.mode = None # generated program vars self._input_vars = defaultdict(list) self._label_vars = defaultdict(list) self._output_vars = defaultdict(list) self._loss_vars = defaultdict(list) self._loss_names = defaultdict(list) self._metric_vars = defaultdict(list) # Consider ProxyLayer as not Paddle inner function because it contains # user-defined layer. for fn_name in [ "_train", "_eval", "_predict", "call_loss", "call_metrics", ]: as_not_paddle_func( f"{inspect.getmodule(ProxyLayer).__name__}.ProxyLayer.{fn_name}" ) @paddle.jit.not_to_static def append_loss_to_shadow_output(self, mode): name = paddle.utils.unique_name.generate('loss') paddle._C_ops.set_persistable_value(self._loss_vars[mode], name) self._loss_names[mode] = name def _train(self, inputs, labels): """ Train process of inner_layer with forward/loss/metric logic. """ # step 1. save feed variables of Program mode = 'train' self._input_vars[mode] = inputs self._label_vars[mode] = labels # step 2. call inner_layer.forward self._output_vars[mode] = self.inner_layer(*inputs) # step 3. calculate loss if needed new_inputs = self._prepare(self.output_vars, labels) self._loss_vars[mode] = self.call_loss(new_inputs) if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ "FLAGS_enable_pir_api" ]: self.append_loss_to_shadow_output(mode) # step 4. calculate metrics if needed self._metric_vars[mode] = self.call_metrics(new_inputs) def _eval(self, inputs, labels): """ Evaluate process of inner_layer with forward/loss/metric logic. """ # TODO(dev): we can reuse codes with self._train after making # sure if they can. # step 1. save feed variables of Program mode = 'eval' self._input_vars[mode] = inputs self._label_vars[mode] = labels # step 2. call inner_layer.forward self._output_vars[mode] = self.inner_layer(*inputs) # step 3. calculate loss if needed new_inputs = self._prepare(self.output_vars, labels) self._loss_vars[mode] = self.call_loss(new_inputs) if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ "FLAGS_enable_pir_api" ]: self.append_loss_to_shadow_output(mode) # step 4. calculate metrics if needed self._metric_vars[mode] = self.call_metrics(new_inputs) def _predict(self, inputs, labels): """ Predict process of inner_layer with forward logic. """ # step 1. save feed variables of Program mode = 'predict' self._input_vars[mode] = inputs self._label_vars[mode] = labels # step 2. call inner_layer.forward self._output_vars[mode] = self.inner_layer(*inputs) @not_to_static def _prepare(self, outputs, labels): """ Concat outputs and labels as a single list NOTE(dev): We use @not_to_static to avoid AST Analysis. """ return to_list(outputs) + to_list(labels) def call_loss(self, inputs): """ Apply Loss Function on outputs and labels. Args: inputs: List[Variable] Returns: List[Variable] """ res = [] if self.loss_func is not None: res = self.loss_func(*inputs) return res def call_metrics(self, inputs): """ Apply Metrics Function on outputs and labels. Args: inputs: List[Variable] Returns: List[Variable] """ outs = [] for metric in self.metrics: outs.append(to_list(metric.compute(*inputs))) return outs def set_mode(self, mode): self.mode = mode self.training = mode == 'train' def clone(self): return ProxyLayer(self.inner_layer, self.loss_func, self.metrics) @property def input_vars(self): return self._input_vars[self.mode] @property def label_vars(self): return self._label_vars[self.mode] @property def output_vars(self): return self._output_vars[self.mode] @property def loss_vars(self): return self._loss_vars[self.mode] @property def loss_names(self): return self._loss_names[self.mode] @property def metric_vars(self): return self._metric_vars[self.mode] @property def startup_program(self): return self.inner_layer._startup_program() class BuildInfo: def __init__(self): self.clear() def has_cache(self, mode, update=False): is_cache = self.states[mode] if update: self.cache(mode) return is_cache def cache(self, mode): self.states[mode] = True def clear(self): self.states = defaultdict(bool) class ProgramHelper: """ A Helper class for Engine to provides different Program IR according specified 'mode'. """ def __init__(self, layer, loss_func, metrics, inputs_spec, labels_spec): # original model config information # TODO(Aurelius84): Implement append_backward and optimizer in ProxyLayer # after distribute engine satisfy basic condition. self.proxy_layer = ProxyLayer(layer, loss_func, metrics) self.inputs_spec = inputs_spec self.labels_spec = labels_spec self.build_info = BuildInfo() self._logger = get_logger(logging.INFO) self.lazy_init = False self._all_params_dist_attr = {} def reset(self): """ Reset all state of current Object. """ self.build_info.clear() self.proxy_layer = self.proxy_layer.clone() def build_program(self, mode): """ Convert dygraph model into static Program IR. """ assert mode in ['train', 'eval', 'predict'] self.proxy_layer.set_mode(mode) # skip if we has already built program. if self.build_info.has_cache(mode, True): self._logger.info( f"Already build program with mode = {mode}, use cached program." ) return self._logger.info(f"start to build program for mode = {mode}.") input_spec = [self.inputs_spec, self.labels_spec] static_func = to_static( self.static_func(), input_spec=input_spec, full_graph=True ) func_name = '_' + mode setattr(self.proxy_layer, func_name, static_func) # NOTE(dev): Because @to_static is a Lazy mechanism, so we explicitly call this to trigger # generating Program IR immediately. concrete_program = getattr(self.proxy_layer, func_name).concrete_program # TODO(zhiqiu): prepare_op_amp_options is not supported for PIR program # It will to use dynamic-static unified amp in pir program, and there is # no need to fit for prepare_op_amp_options if not paddle.base.framework.get_flags("FLAGS_enable_pir_api")[ "FLAGS_enable_pir_api" ]: prepare_op_amp_options( concrete_program.main_program, ProgramTranslator.get_instance()._amp_records, DEFAULT_AMP_OPTIONS, ) self._build_startup_program() def _build_startup_program(self): """ Create and Sync parameters into startup program. """ startup_program = self.startup_program if len(startup_program.global_block().ops) > 1: self.lazy_init = True return for param in self.concrete_program.parameters: Parameter( name=param.name, desc=param, type=param.type, shape=param.shape, dtype=param.dtype, stop_gradient=param.stop_gradient, block=startup_program.global_block(), ) def apply_optimizer(self, optimizer): """ Append backward and generate optimizer operations. """ self._verify_optimizer(optimizer) self._logger.info( "start to apply optimizer: %s ", type(optimizer).__name__ ) # clear optimizer parameters original_params = optimizer._parameter_list optimizer._parameter_list = None with program_guard(self.main_program, self.startup_program): res = optimizer.minimize(self.loss_vars[0]) # restore optimizer parameters optimizer._parameter_list = original_params return res def _verify_optimizer(self, optimizer): assert optimizer is not None assert hasattr(optimizer, "minimize"), ( "Optimizer must have minimize() method." ) assert self.proxy_layer.mode == 'train', ( f"Required mode == 'train', but received '{self.proxy_layer.mode}'" ) assert len(self.loss_vars) == 1, ( f"Required len(loss_vars) == 1, but received len(loss_vars) = {len(self.loss_vars)}" ) def to(self, mode): """ Switch underly proxy layer mode into target mode. """ assert mode in ['train', 'eval', 'predict'] func = getattr(self.proxy_layer, '_' + mode) assert isinstance(func, StaticFunction), ( "Please call build_program(mode) firstly." ) self.proxy_layer.set_mode(mode) def static_func(self): """ Return StaticFunction instance with underly target mode. """ assert self.proxy_layer.mode in [ 'train', 'eval', 'predict', ], "Please call build_program(mode) firstly." func_name = '_' + self.proxy_layer.mode return getattr(self.proxy_layer, func_name) def init_pir(self, main_program, place): # collect all params in current dist program param_values = main_program.global_block().all_parameters() value_name_to_value = {} dy_param_name_to_pir_param_name = {} for value in param_values: value_name_to_value[value.name] = value dy_params = self.concrete_program.parameters[0] pir_param = self.concrete_program.parameters[1] for i in range(len(pir_param)): if pir_param[i].name in value_name_to_value: dy_param_name_to_pir_param_name[dy_params[i].name] = pir_param[ i ].name is_comm = False for param in dy_params: if param.is_dist(): process_mesh, dims_mapping = self._all_params_dist_attr[ param.name ] var_dist_attr = TensorDistAttr() var_dist_attr.process_mesh = process_mesh var_dist_attr.dims_mapping = dims_mapping is_comm = True with paddle.no_grad(): tmp = paddle.base.core.reshard(param, var_dist_attr) if tmp._is_initialized(): param.get_tensor()._share_data_with(tmp.get_tensor()) else: # Only setting the "param" to "None" can't release the memory param.get_tensor()._clear() param = None # create var in scope and share parameters to scope if param is None: continue if param.name not in dy_param_name_to_pir_param_name: # Release the redundant params param.get_tensor()._clear() continue if not param._is_initialized(): continue if param.is_dense(): value_name = dy_param_name_to_pir_param_name[param.name] value = value_name_to_value[value_name] # get param_var's dist_attr assert value.is_dist_dense_tensor_type(), ( f"param [{value.name}] is not dist tensor type" ) dist_attr = { "dims_mapping": value.dist_attr().dims_mapping, "process_shape": value.dist_attr().process_mesh.shape, "process_group": value.dist_attr().process_mesh.process_ids, } # slice param_value with dist_attr # share sliced_param_value with param_tensor in global_scope pir_scope_param = global_scope().var(value_name).get_tensor() sliced_param = Converter.slice_with_dist_attr( param.numpy(), dist_attr ) pir_scope_param.set(sliced_param, place) param.get_tensor()._clear() elif param.is_dist(): value_name = dy_param_name_to_pir_param_name[param.name] value = value_name_to_value[value_name] # assert value.is_dist_dense_tensor_type(), "param [{}] is not dist tensor type".format(value.name) pir_scope_param = global_scope().var(value_name).get_tensor() pir_scope_param._share_data_with( param.get_tensor().get_tensor() ) param.get_tensor()._clear() world_group = get_world_process_group() if ( is_comm and world_group.nranks > 1 and paddle.distributed.get_world_size() > 1 ): paddle.disable_static() barrier_tensor = paddle.full([1], 1, dtype="int32") # barrier is not available in xpu for now if not paddle.framework.core.is_compiled_with_xpu(): paddle._legacy_C_ops.barrier( barrier_tensor, barrier_tensor, 'ring_id', 0 ) paddle.enable_static() def init(self, main_program, place, dist_context): if self.lazy_init: return amp_strategy = dist_context.strategy.amp amp_config = copy.deepcopy(amp_strategy.to_dict()) need_cast_parameter = amp_strategy.enable and amp_config["level"] in [ "o2", "o3", ] is_comm = False for param in self.concrete_program.parameters: if param.is_dist(): serial_main_program = self.concrete_program.main_program var = serial_main_program.global_block().vars[param.name] var_dist_attr = dist_context.get_tensor_dist_attr_for_program( var ) is_comm = True # No need to construct backward. with paddle.no_grad(): tmp = paddle.base.core.reshard(param, var_dist_attr) if tmp._is_initialized(): param.get_tensor()._share_data_with(tmp.get_tensor()) else: # Only setting the "param" to "None" can't release the memory param.get_tensor()._clear() param = None paddle.device.synchronize() # create var in scope and share parameters to scope if param is None: continue if param.name not in main_program.global_block().vars: # Release the redundant params param.get_tensor()._clear() continue if not param._is_initialized(): continue if param.is_dense(): # get param_var's dist_attr var = main_program.global_block().vars[param.name] 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.shape, "process_group": var_dist_attr.process_mesh.process_ids, } # slice param_value with dist_attr # share sliced_param_value with param_tensor in global_scope param_tensor = global_scope().var(param.name).get_tensor() sliced_param = Converter.slice_with_dist_attr( param.numpy(), dist_attr ) param_tensor.set(sliced_param, place) if not need_cast_parameter: param.get_tensor()._clear() elif param.is_dist(): dense_tensor = global_scope().var(param.name).get_tensor() dense_tensor._share_data_with(param.get_tensor().get_tensor()) # transform the parameter in eager mode for amp. if need_cast_parameter: for param in self.concrete_program.parameters: amp_dtype = amp_config["dtype"] scope_var = global_scope().find_var(param.name) # The parameter is not in this rank. if not scope_var: continue # The parameter do not need to transform if param.dtype in [paddle.float16, paddle.bfloat16]: continue scope_tensor = global_scope().var(param.name).get_tensor() assert scope_var and scope_tensor._is_initialized(), ( f"Parameter: {param.name} is not put into global_scope or not initialized." ) param_used = param # For the params without dist_attr. # NOTE(lizhiyu): In principle, each param should have dist_attr. if param.is_dense(): # get param_var's dist_attr var = main_program.global_block().vars[param.name] 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.shape, "process_group": var_dist_attr.process_mesh.process_ids, } # slice param_value with dist_attr sliced_param = Converter.slice_with_dist_attr( param.numpy(), dist_attr ) with paddle.base.dygraph.guard(): param_used = paddle.to_tensor( sliced_param, place=param.place ) param.get_tensor()._clear() with paddle.base.dygraph.guard(): if amp_dtype == "float16": with ( paddle.no_grad(), paddle.base.framework._dygraph_place_guard( place=place ), ): t_casted = param_used.cast( dtype=core.VarDesc.VarType.FP16 ) elif amp_dtype == "bfloat16": with ( paddle.no_grad(), paddle.base.framework._dygraph_place_guard( place=place ), ): t_casted = param_used.cast( dtype=core.VarDesc.VarType.BF16 ) # NOTE(lizhiyu): Clear the origin param. Don't use `param_used.get_tensor().get_tensor()._clear()` to # clear the `DistTensor`, because it can't clear the `_holder`, # which `param_used.get_tensor().get_tensor()` will copy one `DenseTensor`. param_used.get_tensor()._clear() if t_casted.is_dist(): scope_tensor._share_data_with( t_casted.get_tensor().get_tensor() ) else: scope_tensor._share_data_with(t_casted.get_tensor()) world_group = get_world_process_group() if ( is_comm and world_group.nranks > 1 and paddle.distributed.get_world_size() > 1 ): paddle.disable_static() barrier_tensor = paddle.full([1], 1, dtype="int32") # barrier is not available in xpu for now if not paddle.framework.core.is_compiled_with_xpu(): paddle._legacy_C_ops.barrier( barrier_tensor, barrier_tensor, 'ring_id', 0 ) paddle.enable_static() def cache_whole_graph_dist_attr(self, all_params): for param_value in all_params: dist_attr = param_value.dist_attr() if dist_attr: process_mesh = dist_attr.process_mesh dims_mapping = dist_attr.dims_mapping self._all_params_dist_attr[param_value.name] = [ process_mesh, dims_mapping, ] @property def concrete_program(self): return self.static_func().concrete_program @property def main_program(self): return self.concrete_program.main_program @property def startup_program(self): try: return self.proxy_layer.startup_program except Exception as err: self._logger.warning( "The startup_program is not built by `lazy init`." ) if isinstance(err, AssertionError): return self.concrete_program.startup_program raise err @property def input_vars(self): return to_list(self.proxy_layer.input_vars) @property def output_vars(self): return to_list(self.proxy_layer.output_vars) @property def label_vars(self): return to_list(self.proxy_layer.label_vars) @property def loss_vars(self): return to_list(self.proxy_layer.loss_vars) @property def loss_names(self): return to_list(self.proxy_layer.loss_names) @property def metric_vars(self): return to_list(self.proxy_layer.metric_vars) def named_parameters(self): static_func = self.static_func() partial_program = static_func.get_concrete_program( self.inputs_spec, self.labels_spec )[-1] # TODO(xiongkun): support pir in the feature. return {param.name: param for param in partial_program._params}