# 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 paddle logger = logging.getLogger(__name__) formatter = logging.Formatter( fmt='%(asctime)s %(levelname)-8s %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) ch = logging.StreamHandler() ch.setFormatter(formatter) logger.addHandler(ch) from paddle.base import core from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY from paddle.static import Parameter _supported_optimizer_type = [ "adam", "adamax", "adamw", "decayed_adagrad", "momentum", "dgc_momentum", "lars_momentum", "merged_momentum", "lamb", "sgd", ] def tensor_parallel_sync_filter_fn( param, pos_emb=True, layer_norm=True, bias=True ): """ Layer filter function for tensor parallelism transformer. In tensor parallelism of transformer like model, there is 4 kind of param that are supposed to be the same in all tensor parallel peers: * position embedding * scale of layer norm * bias of layer norm * bias of row parallel linear set corresponding input args to select specific layers. NOTE adopting the param name pattern for different transformer blocks. """ p_name = param.name if pos_emb and p_name.startswith("pos_embedding"): return True elif layer_norm and p_name.endswith("_layer_norm_bias"): return True elif layer_norm and p_name.endswith("_layer_norm_scale"): return True elif bias and ".b_" in p_name and (param.is_distributed is False): return True else: return False def resolute_tensor_parallel_ring_id(program): ops = program.global_block().ops ring_id = None for op in ops: if op.type == "c_identity": if ring_id is None: ring_id = int(op.attr("ring_id")) else: assert ring_id == int(op.attr("ring_id")), ( "Found two different ring_id for Tensor Parallel: ring_id={} and ring_id={}.".format( ring_id, int(op.attr("ring_id")) ) ) assert ring_id is not None, "Could NOT found ring_id for Tensor Parallel." return ring_id def copy_parameters(block_, params): for param in params: new_p = Parameter( block=block_, shape=param.shape, dtype=param.dtype, type=param.type, lod_level=( param.lod_level if param.type == core.VarDesc.VarType.DENSE_TENSOR else None ), stop_gradient=param.stop_gradient, trainable=param.trainable, optimize_attr=param.optimize_attr, regularizer=param.regularizer, error_clip=param.error_clip, name=param.name, ) assert param.is_distributed is False, ( f"Try to sync Distributed Parameter: {param}" ) new_p.is_distributed = False block_.vars[new_p.name] = new_p def insert_sync_op( block, idx, tp_degree, sync_mode, sync_ring_id, src_rank, varname, op_role ): if sync_mode == "broadcast": block._insert_op_without_sync( idx, type='broadcast', inputs={'x': varname}, outputs={'out': varname}, attrs={ 'ring_id': sync_ring_id, 'root': src_rank, OP_ROLE_KEY: op_role, }, ) elif sync_mode == "average": block._insert_op_without_sync( idx, type='scale', inputs={'X': varname}, outputs={'Out': varname}, attrs={'scale': 1.0 / tp_degree, OP_ROLE_KEY: op_role}, ) block._insert_op_without_sync( idx, type='all_reduce', inputs={'x': varname}, outputs={'out': varname}, attrs={ 'ring_id': sync_ring_id, 'reduce_type': paddle.distributed.ReduceOp.SUM, OP_ROLE_KEY: op_role, }, ) else: raise NotImplementedError( f'Sync mode of [{sync_mode}] is NOT supported.' ) def insert_synchronization( block, params_to_sync, tp_degree, sync_ring_id, sync_param, sync_grad, sync_moment, sync_mode, src_rank, ): unsync_param_names = [p.name for p in params_to_sync] for idx, op in reversed(list(enumerate(block.ops))): if op.type in _supported_optimizer_type: assert "Param" in op.input_names assert len(op.input("Param")) == 1 param_name = op.input("Param")[0] op_role = op.attr(OP_ROLE_KEY) if param_name in unsync_param_names: unsync_param_names.remove(param_name) # Param sync after opt if sync_param: assert ( "ParamOut" in op.output_names and op.output("ParamOut")[0] == param_name ) insert_sync_op( block, idx + 1, tp_degree, sync_mode, sync_ring_id, src_rank, param_name, op_role, ) if ( "MasterParamOut" in op.output_names and len(op.output("MasterParamOut")) == 1 ): sync_var = op.output("MasterParamOut")[0] insert_sync_op( block, idx + 1, tp_degree, sync_mode, sync_ring_id, src_rank, sync_var, op_role, ) # Moment sync after opt if sync_moment: if ( "Moment1Out" in op.output_names and len(op.output("Moment1Out")) == 1 ): sync_var = op.output("Moment1Out")[0] insert_sync_op( block, idx + 1, tp_degree, sync_mode, sync_ring_id, src_rank, sync_var, op_role, ) if ( "Moment2Out" in op.output_names and len(op.output("Moment2Out")) == 1 ): sync_var = op.output("Moment2Out")[0] insert_sync_op( block, idx + 1, tp_degree, sync_mode, sync_ring_id, src_rank, sync_var, op_role, ) # Grad sync before opt if sync_grad: assert ( "Grad" in op.input_names and len(op.input("Grad")) == 1 ) sync_var = op.input("Grad")[0] insert_sync_op( block, idx, tp_degree, sync_mode, sync_ring_id, src_rank, sync_var, op_role, ) assert len(unsync_param_names) == 0, ( f"The following param is unsync by some error: {unsync_param_names}" ) def add_extra_synchronization( program, params_filter_fn=tensor_parallel_sync_filter_fn, tp_degree=8, sync_mode="broadcast", sync_param=True, sync_grad=False, sync_moment=False, src_rank=0, sync_ring_id=None, ): """ Inplace add extra synchronization for input program. program(Paddle.Program): distributed train program. params_filter_fn(callable): function to filter out parameter for synchronization. sync_mode(string): select from "broadcast": parameter is sync by broadcasted from 'src_rank' to all other ranks. "average": parameter is sync by average among all ranks src_rank(int): the src used in broadcast sync_mode. sync_param(bool): extra synchronize parameters. sync_grad(bool): extra synchronize gradients. sync_grad(bool): extra synchronize optimizer momentum. sync_ring_id(int): communicator id use for synchronization, if it is None, use the ring_id of tensor parallel. """ logger.info("Constructing Extra Parameter Synchronization.") logger.info( f"Tensor Parallel Degree: {tp_degree}, Synchronization mode: {sync_mode}" ) # adopt for pipeline opt if program._pipeline_opt is not None: assert program._pipeline_opt['section_program'] is not None, ( "Pipeline is enable but section_program is None" ) program = program._pipeline_opt['section_program'] # step1: collect the param that need to be sync params_to_sync = [] # TODO support multiple blocks with different parameter. all_params = program.global_block().all_parameters() for param in all_params: if params_filter_fn(param): params_to_sync.append(param) logger.info( "The following param are going to be synchronization everytime the optimizer update phase of the program is run: " ) logger.info([p.name for p in params_to_sync]) # step2: resolute synchronization communicator group (ring_id) if sync_ring_id is None: sync_ring_id = resolute_tensor_parallel_ring_id(program) # step3: insert synchronization # TODO support gradient merge with different update block block = program.global_block() insert_synchronization( block, params_to_sync, tp_degree, sync_ring_id, sync_param, sync_grad, sync_moment, sync_mode, src_rank, )