# Copyright (c) 2023 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 logging import numpy as np import paddle from paddle.distributed.auto_parallel.static.utils import ( is_optimize_op, is_recompute_op, naive_set_dist_op_attr_for_program_by_mesh_and_mapping, set_var_dist_attr, ) from paddle.utils import unique_name from ..utils.log_utils import get_logger from .auto_parallel_sharding import ( _inference_data_parallel_group_for_operator, _is_reshard_op, _skip_ops, is_forward_op, ) from .pass_base import PassBase, register_pass logger = get_logger(logging.INFO, "FusedLinearPromotionPass") _supported_optimizer_type = [ "adam", "adamax", "adamw", "decayed_adagrad", "momentum", "dgc_momentum", "lars_momentum", "merged_momentum", "lamb", "sgd", ] FUSED_LINEAR_SOURCE_PATTERNS_LIST = [ # amp_level == 'o2' or 'o3' { # only MP "forward": ["matmul_v2", "all_reduce", "elementwise_add"], "backward": ["elementwise_add_grad", "matmul_v2_grad"], }, { # MP + SP "forward": ["matmul_v2", "reduce_scatter", "elementwise_add"], "backward": [ "elementwise_add_grad", "all_reduce", "scale", "all_gather", "matmul_v2_grad", "all_gather", ], }, { # DP + MP "forward": ["matmul_v2", "all_reduce", "elementwise_add"], "backward": [ "elementwise_add_grad", "all_reduce", "scale", "matmul_v2_grad", ], }, { # DP + MP + SP "forward": ["matmul_v2", "reduce_scatter", "elementwise_add"], "backward": [ "elementwise_add_grad", "all_reduce", "scale", "all_reduce", "scale", "all_gather", "matmul_v2_grad", "all_gather", ], }, # amp_level == 'o1' { "forward": ["matmul_v2", "all_reduce", "cast", "elementwise_add"], "backward": ["elementwise_add_grad", "matmul_v2_grad"], }, { "forward": ["matmul_v2", "reduce_scatter", "cast", "elementwise_add"], "backward": [ "elementwise_add_grad", "all_reduce", "scale", "all_gather", "all_gather", "matmul_v2_grad", ], }, { "forward": ["matmul_v2", "all_reduce", "cast", "elementwise_add"], "backward": [ "elementwise_add_grad", "all_reduce", "scale", "matmul_v2_grad", ], }, { "forward": ["matmul_v2", "reduce_scatter", "cast", "elementwise_add"], "backward": [ "elementwise_add_grad", "all_reduce", "scale", "all_reduce", "scale", "all_gather", "matmul_v2_grad", "all_gather", ], }, ] @register_pass("auto_parallel_fused_linear_promotion") class FusedLinearPromotionPass(PassBase): """ Apply pre-promotion that specialized for fused_linear_pass in tensor parallelism or sequence parallelism in Auto Parallel. """ def __init__(self): super().__init__() self.set_attr("dist_context", None) self.set_attr("global_rank", -1) self.set_attr("enable_sp", False) self.set_attr("amp_level", "o0") self.set_attr("params_grads", None) def _check_self(self): if self.get_attr("dist_context") is None: return False if (not isinstance(self.get_attr("global_rank"), int)) or self.get_attr( "global_rank" ) < 0: return False return True def _check_conflict(self, other_pass): return True def _apply_single_impl(self, main_program, startup_program, context): self._dist_context = self.get_attr("dist_context") self._global_rank = int(self.get_attr("global_rank")) self._params_grads = self.get_attr("params_grads") self._amp_level = self.get_attr("amp_level") self._enable_sp = self.get_attr("enable_sp") self._is_amp_o1 = self._amp_level == 'o1' self._source_patterns = {} self._enable_dp, self._enable_mp = self._is_enable_dp_mp( self._dist_context ) pattern_offset = 4 if self._is_amp_o1 else 0 if self._enable_sp: if self._enable_dp: self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[ 3 + pattern_offset ] else: self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[ 1 + pattern_offset ] elif self._enable_mp: if self._enable_dp: self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[ 2 + pattern_offset ] else: self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[ 0 + pattern_offset ] else: logger.warning("Neither of sp and mp is enabled, skip this pass") return dp_group = None if self._enable_dp: dp_group = self._collective_data_parallel_groups( main_program.global_block() ) # 1. get whether the current rank is first rank in mp self._is_first_rank = self._is_tp_sp_first_rank( self._dist_context, self._global_rank ) logger.debug(f"before main_program: {main_program}") # 2. get the forward and backward op list indexes in source patterns ( forward_segments, backward_segments, ) = self._get_forward_backward_op_segments(main_program) if len(forward_segments) == 0 or len(backward_segments) == 0: logger.warning( "No forward and backward op segments, skip this pass" ) return # 3 transform the forward ops rename_var_names_map, deleted_bias_names = self._transform_forward( main_program, forward_segments, backward_segments, self._is_first_rank, self._enable_sp, self._is_amp_o1, ) # 4 transform the backward ops self._transform_backward( main_program, backward_segments, rename_var_names_map, self._is_first_rank, self._enable_sp, ) # 5. transform the optimizer ops self._transform_opt( main_program, deleted_bias_names, self._params_grads, self._is_first_rank, self._is_amp_o1, ) logger.info(f"deleted_bias_names: {deleted_bias_names}") logger.debug(f"after main_program: {main_program}") # 6. transform the startup program self._transform_startup_program( startup_program, deleted_bias_names, dp_group, self._is_first_rank ) def _is_tp_sp_first_rank(self, dist_context, rank): for process_mesh in dist_context.process_meshes: inner_mesh_shape = process_mesh.shape inner_mesh = (np.array(process_mesh.process_ids)).reshape( inner_mesh_shape ) if len(inner_mesh_shape) == 1: return rank == min(process_mesh.process_ids) elif len(inner_mesh.shape) == 2: for id0 in range(inner_mesh_shape[0]): if rank == min(inner_mesh[id0, :]): return True elif len(inner_mesh.shape) == 3: for id0 in range(inner_mesh_shape[0]): for id1 in range(inner_mesh_shape[1]): if rank == min(inner_mesh[id0, id1, :]): return True else: raise ValueError("inner mesh shape is not supported") return False def _is_enable_dp_mp(self, dist_context): for process_mesh in dist_context.process_meshes: inner_mesh_shape = process_mesh.shape inner_mesh = (np.array(process_mesh.process_ids)).reshape( inner_mesh_shape ) if len(inner_mesh_shape) == 1: return False, inner_mesh_shape[0] > 1 else: # DP * MP return inner_mesh_shape[-2] > 1, inner_mesh_shape[-1] > 1 return False, False def _get_forward_backward_op_segments(self, main_program): """ Get the operator segments according to the source patterns. """ def can_match_pattern( ops, start_id, pattern, forward_matmul_inputs, is_backward=False ): """ Check whether the ops in the range [start_id, start_id + len(pattern)] can match the pattern. If the ops is in forward pass, check it directly. However, when the ops is in backward pass, we need to additionally check whether the input of the last op in pattern is in forward_matmul_inputs to deal the case of enabling recompute. """ new_id = start_id if not is_backward: for op_name in pattern: if ops[new_id].type != op_name: return False new_id += 1 forward_matmul_inputs.extend(ops[start_id].input_arg_names) return True else: for op_name in pattern: if ops[new_id].type != op_name: return False new_id += 1 matmul_grad_input_names = ops[new_id - 1].input_arg_names # for refined-recompute if ( matmul_grad_input_names[1] not in forward_matmul_inputs and matmul_grad_input_names[2] not in forward_matmul_inputs ): return False return True global_block = main_program.global_block() forward_segments = [] backward_segments = [] ops_len = len(global_block.ops) self._forward_patterns_len = len(self._source_patterns["forward"]) self._backward_patterns_len = len(self._source_patterns["backward"]) forward_matmul_inputs = [] for id, op in enumerate(global_block.ops): if id > ops_len - self._backward_patterns_len: break if int(op.desc.attr('op_role')) == 0 or ( is_recompute_op(op) and not op.type.endswith("_grad") ): # forward if can_match_pattern( global_block.ops, id, self._source_patterns["forward"], forward_matmul_inputs, is_backward=False, ): forward_segments.append( [id, id + self._forward_patterns_len] ) elif int(op.desc.attr('op_role')) == 1: # backward if can_match_pattern( global_block.ops, id, self._source_patterns["backward"], forward_matmul_inputs, is_backward=True, ): backward_segments.append( [id, id + self._backward_patterns_len] ) else: pass assert len(forward_segments) >= len(backward_segments), ( "The number of forward segments should be not shorter than the number of backward segments." ) logger.info(f"forward_segments: {forward_segments}") logger.info(f"backward_segments: {backward_segments}") return forward_segments, backward_segments def _collective_data_parallel_groups(self, main_block): for op in main_block.ops: if not is_forward_op(op) or op.type in _skip_ops: continue # NOTE: there aren't dist_attr in the ops which reshard insert, # and should be skip in sharding. if _is_reshard_op(op): continue group = _inference_data_parallel_group_for_operator( self._global_rank, op, self._dist_context ) if group is not None: return group return None def _transform_forward( self, main_program, forward_segments, backward_segments, is_first_rank, is_sp, is_amp_o1, ): """ Transform the forward pass. """ def _transform_forward_segment( global_block, forward_segment, backward_segments, is_first_rank, is_sp, is_amp_o1, ): """ Transform one forward segment. """ # 1. prepare the forward_segment # 1.1 check whether the forward_segment is right origin_matmul_op = global_block.ops[forward_segment[0]] origin_comm_op = global_block.ops[forward_segment[0] + 1] origin_add_op = global_block.ops[forward_segment[1] - 1] origin_cast_op = ( global_block.ops[forward_segment[1] - 2] if is_amp_o1 else None ) origin_matmul_output_name = origin_matmul_op.output_arg_names[0] origin_comm_input_name = origin_comm_op.input_arg_names[0] assert origin_matmul_output_name == origin_comm_input_name, ( f"The 0th op output name {origin_matmul_output_name} is not equal to the 1st op input name {origin_comm_input_name}" ) origin_comm_output_name = origin_comm_op.output_arg_names[0] origin_add_input_names = origin_add_op.input_arg_names assert origin_comm_output_name == origin_add_input_names[0], ( f"The 1st op output name {origin_comm_output_name} is not equal to the 2nd op input name {origin_add_input_names[0]}" ) # 1.2 get the origin dist_attr origin_add_dist_attr = ( self._dist_context.get_op_dist_attr_for_program(origin_add_op) ) assert origin_add_dist_attr is not None, ( f"Origin add op {origin_add_op.type} has no dist attr" ) ref_mesh = origin_add_dist_attr.process_mesh in_var_dist_attr = origin_add_dist_attr.get_input_dist_attr( origin_add_op.input_arg_names[0] ) ref_mapping = in_var_dist_attr.dims_mapping # 2. deal matmul_v2 op origin_matmul_output_new_name = unique_name.generate( origin_matmul_output_name + "@promote" ) origin_matmul_output_new_var = global_block.create_var( name=origin_matmul_output_new_name, dtype=global_block.var(origin_matmul_output_name).dtype, shape=global_block.var(origin_matmul_output_name).shape, persistable=False, stop_gradient=False, ) set_var_dist_attr( self._dist_context, origin_matmul_output_new_var, ref_mapping, ref_mesh, ) rename_vars_map[origin_matmul_output_name] = ( origin_matmul_output_new_name ) origin_matmul_op._rename_output( origin_matmul_output_name, origin_matmul_output_new_name ) naive_set_dist_op_attr_for_program_by_mesh_and_mapping( origin_matmul_op, ref_mesh, ref_mapping, self._dist_context ) # 3. deal add op and cast op if is_first_rank: # insert the "elementwise_add" op before reduce_sum new_add_op = global_block._insert_op_without_sync( forward_segment[0] + 1, type="nop", ) new_op_desc = new_add_op.desc new_op_desc.copy_from(origin_add_op.desc) # create new var of new_add_op output origin_add_output_name = origin_add_op.output_arg_names[0] new_add_op_output_name = unique_name.generate( origin_add_output_name + "@promote" ) new_shape_var_name = ( origin_add_output_name if not is_sp else origin_matmul_output_name ) global_block.create_var( name=new_add_op_output_name, dtype=global_block.var(origin_add_output_name).dtype, shape=global_block.var(new_shape_var_name).shape, persistable=False, stop_gradient=False, ) global_block._remove_var( origin_matmul_output_name ) # We can remove the origin_matmul_output now. global_block._remove_var(origin_add_output_name) new_add_op._rename_output( origin_add_output_name, new_add_op_output_name ) rename_vars_map[origin_add_op.input_arg_names[0]] = ( origin_matmul_output_new_name ) new_add_op._rename_input( origin_add_op.input_arg_names[0], origin_matmul_output_new_name, ) # deal dist_attr naive_set_dist_op_attr_for_program_by_mesh_and_mapping( new_add_op, ref_mesh, ref_mapping, self._dist_context ) # 'cast' op also need to adjust if is_amp_o1: new_cast_op = global_block._insert_op_without_sync( forward_segment[0] + 1, type="nop", ) new_op_desc = new_cast_op.desc new_op_desc.copy_from(origin_cast_op.desc) if ( new_cast_op.input_arg_names[0] not in delete_bias_vars_name ): # fp16 = cast(fp32) delete_bias_vars_name.append( new_cast_op.input_arg_names[0] ) else: if ( new_add_op.input_arg_names[1] not in delete_bias_vars_name ): delete_bias_vars_name.append( new_add_op.input_arg_names[1] ) else: # We can remove the origin_matmul_output now. origin_add_output_name = origin_add_op.output_arg_names[0] global_block._remove_var(origin_add_output_name) global_block._remove_var(origin_matmul_output_name) # 4. deal comm op # The input of all_reduce_sum only be used once, so we don't need add it in the rename_vars_map if is_first_rank: origin_comm_op._rename_input( origin_comm_op.input_arg_names[0], new_add_op.output_arg_names[0], ) else: origin_comm_op._rename_input( origin_comm_op.input_arg_names[0], origin_matmul_output_new_name, ) if ( origin_comm_op.type == "all_reduce" and origin_comm_op.attr("reduce_type") == paddle.distributed.ReduceOp.SUM ): new_comm_var_name = origin_comm_op.input_arg_names[0] else: new_comm_var_name = unique_name.generate( origin_comm_output_name + "@promote" ) global_block.create_var( name=new_comm_var_name, dtype=global_block.var(origin_comm_output_name).dtype, shape=global_block.var(origin_comm_output_name).shape, persistable=False, stop_gradient=False, ) rename_vars_map[origin_comm_output_name] = new_comm_var_name if global_block.has_var(origin_comm_output_name): global_block._remove_var(origin_comm_output_name) rename_vars_map[origin_add_output_name] = ( new_comm_var_name # the output of comm op inplace the output of add op for next ops ) origin_comm_op._rename_output( origin_comm_output_name, new_comm_var_name ) naive_set_dist_op_attr_for_program_by_mesh_and_mapping( origin_comm_op, ref_mesh, ref_mapping, self._dist_context ) # 5. remove elementwise_add op and cast op if is_first_rank: if is_amp_o1: global_block._remove_op(forward_segment[0] + 5) global_block._remove_op(forward_segment[0] + 4) else: global_block._remove_op(forward_segment[0] + 3) else: global_block._remove_op( forward_segment[1] - 1 ) # remove elementwise_add op if is_amp_o1: if ( origin_cast_op.input_arg_names[0] not in delete_bias_vars_name ): delete_bias_vars_name.append( origin_cast_op.input_arg_names[0] ) global_block._remove_var(origin_cast_op.output_arg_names[0]) global_block._remove_op( forward_segment[1] - 2 ) # remove cast op else: if origin_add_input_names[1] not in delete_bias_vars_name: delete_bias_vars_name.append(origin_add_input_names[1]) # update backward forward_segment for back_seg in reversed(backward_segments): if is_amp_o1: if back_seg[0] > forward_segment[0]: back_seg[0] -= 2 back_seg[1] -= 2 else: break else: if back_seg[0] > forward_segment[0]: back_seg[0] -= 1 back_seg[1] -= 1 else: break global_block = main_program.global_block() rename_vars_map = {} # origin_name -> new_name delete_bias_vars_name = [] for segment in reversed(forward_segments): _transform_forward_segment( global_block, segment, backward_segments, is_first_rank, is_sp, is_amp_o1, ) global_block._sync_with_cpp() return rename_vars_map, delete_bias_vars_name def _transform_backward( self, main_program, backward_segments, rename_var_names_map, is_first_rank, is_sp, ): global_block = main_program.global_block() to_delete_grad_of_param = [] if is_first_rank: if is_sp: # place the comm_op(all_gather) before the elementwise_add_grad for segment in reversed(backward_segments): add_grad_op = global_block.ops[segment[0]] matmul_grad_op = global_block.ops[segment[-1] - 1] origin_comm_op_id = segment[-1] - 2 origin_comm_op = global_block.ops[origin_comm_op_id] new_comm_op = global_block._insert_op( segment[0], type="nop", ) new_comm_op.desc.copy_from(origin_comm_op.desc) # rename input and output new_comm_op._rename_input( origin_comm_op.input_arg_names[0], add_grad_op.input_arg_names[0], ) add_grad_op._rename_input( add_grad_op.input_arg_names[0], new_comm_op.output_arg_names[0], ) matmul_grad_op._rename_input( matmul_grad_op.input_arg_names[0], add_grad_op.output_arg_names[0], ) global_block._remove_op(segment[-1] - 1) if self._enable_dp: global_block._remove_op(segment[0] + 5) # scale global_block._remove_op( segment[0] + 4 ) # all_reduce_sum else: global_block._remove_op(segment[0] + 3) # scale global_block._remove_op( segment[0] + 2 ) # all_reduce_sum global_block._sync_with_cpp() else: # not is_first_rank_in tp or sp # need to delete the grad op associated with the deleted bias var if not is_sp: for segment in reversed(backward_segments): add_grad_op = global_block.ops[segment[0]] rename_var_names_map[add_grad_op.output_arg_names[0]] = ( add_grad_op.input_arg_names[0] ) global_block._remove_var(add_grad_op.output_arg_names[0]) to_delete_grad_of_param.append( add_grad_op.output_arg_names[1] ) if self._enable_dp: global_block._remove_op(segment[0] + 2) # scale op global_block._remove_op( segment[0] + 1 ) # all_reduce_sum op global_block._remove_op(segment[0]) global_block._sync_with_cpp() else: for segment in reversed(backward_segments): add_grad_op = global_block.ops[segment[0]] origin_comm_op = global_block.ops[segment[-1] - 2] rename_var_names_map[add_grad_op.output_arg_names[0]] = ( add_grad_op.input_arg_names[0] ) origin_comm_op._rename_input( origin_comm_op.input_arg_names[0], add_grad_op.input_arg_names[0], ) global_block._remove_var(add_grad_op.output_arg_names[0]) to_delete_grad_of_param.append( add_grad_op.output_arg_names[1] ) if self._enable_dp: # DP global_block._remove_op( segment[0] + 4 ) # scale op for dp global_block._remove_op( segment[0] + 3 ) # all_reduce_sum op for dp global_block._remove_op(segment[0] + 2) # scale op for sp global_block._remove_op( segment[0] + 1 ) # all_reduce_sum op for sp global_block._remove_op( segment[0] ) # elementwise_add_grad op global_block._sync_with_cpp() # rename input vars in global_block for op in global_block.ops: if is_optimize_op(op): continue for var_name in op.input_arg_names: if var_name in rename_var_names_map: op._rename_input(var_name, rename_var_names_map[var_name]) if self._is_amp_o1: for var_name in to_delete_grad_of_param: global_block._remove_var(var_name) global_block._sync_with_cpp() def _transform_opt( self, main_program, deleted_bias_names, params_grads, is_first_rank, is_amp_o1, ): if is_first_rank: return deleted_bias_grads_names = [] to_delete_params_grads = [] for id, (param, grad) in enumerate(params_grads): if param.name in deleted_bias_names: deleted_bias_grads_names.append(grad.name) to_delete_params_grads.append(id) to_delete_op_ids = [] for id in reversed(range(len(main_program.global_block().ops))): global_block = main_program.global_block() op = global_block.ops[id] op_input_names = op.input_arg_names for op_input in op_input_names: if op_input in deleted_bias_grads_names: if op.type in _supported_optimizer_type: for output_var in op.output_arg_names: global_block._remove_var(output_var) grad_var = op.input('Grad')[0] global_block._remove_var(grad_var) to_delete_op_ids.append(id) if ( op.type == "squared_l2_norm" or op.type == "clip_by_norm" ): output_var_name = op.output_arg_names[0] global_block._remove_var(output_var_name) to_delete_op_ids.append(id) for intra_id in range(id + 1, len(global_block.ops)): intra_op = global_block.ops[intra_id] if ( output_var_name in intra_op.input_arg_names and intra_op.type == "stack" ): origin_vars = intra_op.input("X") origin_vars.remove(output_var_name) intra_op.desc.set_input("X", origin_vars) break if op.type == "elementwise_mul": to_delete_op_ids.append(id) # check_finite_and_unscale and update_loss_scaling if ( op.type == "check_finite_and_unscale" or op.type == "update_loss_scaling" ): origin_vars = op.input("X") origin_vars.remove(op_input) op.desc.set_input("X", origin_vars) origin_vars = op.output("Out") origin_vars.remove(op_input) op.desc.set_output("Out", origin_vars) if is_amp_o1: for output_name in op.output_arg_names: if ( output_name in deleted_bias_grads_names and op.type == 'cast' ): to_delete_op_ids.append(id) for id in to_delete_op_ids: global_block._remove_op(id) main_program.global_block()._sync_with_cpp() for id in reversed(to_delete_params_grads): del params_grads[id] return def _transform_startup_program( self, startup_program, deleted_bias_names, dp_group, is_first_rank ): """ Delete the vars and ops associated with deleted_bias_names in startup program. """ logger.debug(f"Before transform startup_program: {startup_program}") cur_glock = startup_program.global_block() to_delete_op_ids = [] # for variables associated with deleted_bias_names in amp-o2, such as 'opt_linear_1.b_0_fp32_master_0' to_delete_extra_vars = [] for id, op in enumerate(cur_glock.ops): if not is_first_rank: output_var = op.output_arg_names[0] if output_var in deleted_bias_names: to_delete_op_ids.append(id) else: for var_name in deleted_bias_names: if var_name in output_var: to_delete_op_ids.append(id) if output_var not in to_delete_extra_vars: to_delete_extra_vars.append(output_var) else: if op.type == "broadcast": input_vars = op.input_arg_names if ( input_vars[0] in deleted_bias_names and id not in to_delete_op_ids ): if dp_group is None or ( dp_group is not None and op.attr("ring_id") != dp_group.id ): to_delete_op_ids.append(id) for to_delete_id in reversed(to_delete_op_ids): cur_glock._remove_op(to_delete_id) if not is_first_rank: for var_name in deleted_bias_names: cur_glock._remove_var(var_name) for var_name in to_delete_extra_vars: if cur_glock.has_var(var_name): cur_glock._remove_var(var_name) cur_glock._sync_with_cpp() logger.debug(f"After transform startup_program: {startup_program}")