# Copyright (c) 2024 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 from collections import OrderedDict import paddle from paddle.base import core from ...auto_parallel.static.utils import OpRole from ...utils.log_utils import get_logger from ..pass_base import register_pass from ..pass_utils import ( _create_program_and_ops, _get_device, _pir_get_backward_op_type, _pir_overlap_send_recv, _pir_split_matmul_grad_to_matmul, infer_chunk_id, ) from .pipeline_pass_base import PipelinePassBase logger = get_logger(logging.INFO) @register_pass("pipeline_scheduler_VPP") class PipelineVirtualPipelinePass(PipelinePassBase): def __init__(self): super().__init__() self._real_overlap_sharding_reduce = False self.reduce_comm_suffix = "_reduce" self._forward_micro_step_counter = {} self._backward_micro_step_counter = {} self.jobs_in_stable_phase_in_pir = [ self.BACKWARD, self.RECV_FORWARD, self.SEND_BACKWARD, self.FORWARD, ] def _record_fwd_micro_step(self, virtual_pp_rank): real_micro_step = self._forward_micro_step_counter[virtual_pp_rank] self._forward_micro_step_counter[virtual_pp_rank] += 1 return real_micro_step def _record_bwd_micro_step(self, virtual_pp_rank): real_micro_step = self._backward_micro_step_counter[virtual_pp_rank] self._backward_micro_step_counter[virtual_pp_rank] += 1 return real_micro_step def _create_job_list(self): if self._in_pir_mode: return self._pir_create_job_list() accumulate_steps = self.get_attr("num_micro_batches") stage_id = self.get_attr("pp_stage") num_stages = self.get_attr("pp_degree") num_model_chunks = self.get_attr("vpp_degree") split_backward = self.get_attr("split_backward", False) remainder = accumulate_steps % num_stages for i in range(num_model_chunks): self._forward_micro_step_counter[i] = 0 self._backward_micro_step_counter[i] = 0 assert accumulate_steps >= num_stages def _get_virtual_pp_rank(micro_step, forward): virtual_pp_stage = micro_step % (num_stages * num_model_chunks) if micro_step <= (accumulate_steps // num_stages) * ( num_stages * num_model_chunks ): virtual_pp_stage = virtual_pp_stage // num_stages else: virtual_pp_stage = virtual_pp_stage // remainder if not forward: virtual_pp_stage = num_model_chunks - virtual_pp_stage - 1 return virtual_pp_stage total_num_steps = accumulate_steps * num_model_chunks if accumulate_steps == num_stages: warmup_steps = total_num_steps else: warmup_steps = (num_stages - stage_id - 1) * 2 warmup_steps += (num_model_chunks - 1) * num_stages warmup_steps = min(warmup_steps, total_num_steps) steady_steps = total_num_steps - warmup_steps real_split_backward = ( accumulate_steps == num_stages ) and split_backward job_list = [] for micro_step in range(warmup_steps): virtual_pp_rank = _get_virtual_pp_rank(micro_step, forward=True) micro_batch_id = self._record_fwd_micro_step(virtual_pp_rank) fw_job = core.Job(self.FORWARD + str(virtual_pp_rank)) fw_job.set_micro_batch_id(micro_batch_id) job_list.append(fw_job) for micro_step in range(steady_steps): fwd_micro_step = micro_step + warmup_steps fwd_virtual_pp_rank = _get_virtual_pp_rank( fwd_micro_step, forward=True ) fwd_micro_batch_id = self._record_fwd_micro_step( fwd_virtual_pp_rank ) fwd_job = core.Job(self.FORWARD + str(fwd_virtual_pp_rank)) fwd_job.set_micro_batch_id(fwd_micro_batch_id) job_list.append(fwd_job) bw_micro_step = micro_step bwd_virtual_pp_rank = _get_virtual_pp_rank( bw_micro_step, forward=False ) bwd_micro_batch_id = self._record_bwd_micro_step( bwd_virtual_pp_rank ) if real_split_backward: bwd_job = core.Job( self.BACKWARD + "_b" + str(bwd_virtual_pp_rank) ) else: bwd_job = core.Job(self.BACKWARD + str(bwd_virtual_pp_rank)) bwd_job.set_micro_batch_id(bwd_micro_batch_id) job_list.append(bwd_job) for micro_step in range(steady_steps, total_num_steps): virtual_pp_rank = _get_virtual_pp_rank(micro_step, forward=False) micro_batch_id = self._record_bwd_micro_step(virtual_pp_rank) if real_split_backward: bwd_job = core.Job(self.BACKWARD + "_b" + str(virtual_pp_rank)) else: bwd_job = core.Job(self.BACKWARD + str(virtual_pp_rank)) bwd_job.set_micro_batch_id(micro_batch_id) job_list.append(bwd_job) # TODO(lizhiyu): Inserting 'backward_b' and 'backward_w' interleavedly can decrease the memory, # but it reduces the speed. We should find the better way to use the code here. # next_virtual_pp_rank = _get_virtual_pp_rank(micro_step + 1, forward=False) # if next_virtual_pp_rank != virtual_pp_rank: # for micro_batch_id in range(0, accumulate_steps): # w_job = core.Job(BACKWARD + "_w" + str(virtual_pp_rank)) # w_job.set_micro_batch_id(micro_batch_id) # job_list.append(w_job) if real_split_backward: for chunk_id in range(num_model_chunks - 1, -1, -1): for micro_batch_id in range(0, accumulate_steps): if ( self._real_overlap_sharding_reduce and micro_batch_id == accumulate_steps - 1 ): w_job = core.Job( self.BACKWARD + "_w" + str(chunk_id) + self.reduce_comm_suffix ) else: w_job = core.Job(self.BACKWARD + "_w" + str(chunk_id)) w_job.set_micro_batch_id(micro_batch_id) job_list.append(w_job) job_types = [job.type() for job in job_list] logger.debug(f"The VPP job list: {job_types}") opt_job = core.Job(self.OPT) job_list.append(opt_job) return job_list def _pir_create_job_list(self): accumulate_steps = self.get_attr("num_micro_batches") stage_id = self.get_attr("pp_stage") num_stages = self.get_attr("pp_degree") num_model_chunks = self.get_attr("vpp_degree") split_backward = self.get_attr("split_backward", False) remainder = accumulate_steps % num_stages for i in range(num_model_chunks): self._forward_micro_step_counter[i] = 0 self._backward_micro_step_counter[i] = 0 assert accumulate_steps >= num_stages def _get_virtual_pp_rank(micro_step, forward): virtual_pp_stage = micro_step % (num_stages * num_model_chunks) if micro_step <= (accumulate_steps // num_stages) * ( num_stages * num_model_chunks ): virtual_pp_stage = virtual_pp_stage // num_stages else: virtual_pp_stage = virtual_pp_stage // remainder if not forward: virtual_pp_stage = num_model_chunks - virtual_pp_stage - 1 return virtual_pp_stage total_num_steps = accumulate_steps * num_model_chunks if accumulate_steps == num_stages: warmup_steps = total_num_steps else: warmup_steps = (num_stages - stage_id - 1) * 2 warmup_steps += (num_model_chunks - 1) * num_stages warmup_steps = min(warmup_steps, total_num_steps) real_split_backward = ( accumulate_steps == num_stages ) and split_backward if not real_split_backward: warmup_steps = min(total_num_steps, warmup_steps + 1) steady_steps = total_num_steps - warmup_steps job_list = [] for micro_step in range(warmup_steps): virtual_pp_rank = _get_virtual_pp_rank(micro_step, forward=True) micro_batch_id = self._record_fwd_micro_step(virtual_pp_rank) if not real_split_backward: recv_fwd_job = core.Job( self.RECV_FORWARD + str(virtual_pp_rank) ) recv_fwd_job.set_micro_batch_id(micro_batch_id) job_list.append(recv_fwd_job) fw_job = core.Job(self.FORWARD + str(virtual_pp_rank)) fw_job.set_micro_batch_id(micro_batch_id) job_list.append(fw_job) if real_split_backward: for micro_step in range(steady_steps): fwd_micro_step = micro_step + warmup_steps fwd_virtual_pp_rank = _get_virtual_pp_rank( fwd_micro_step, forward=True ) fwd_micro_batch_id = self._record_fwd_micro_step( fwd_virtual_pp_rank ) fwd_job = core.Job(self.FORWARD + str(fwd_virtual_pp_rank)) fwd_job.set_micro_batch_id(fwd_micro_batch_id) job_list.append(fwd_job) bw_micro_step = micro_step bwd_virtual_pp_rank = _get_virtual_pp_rank( bw_micro_step, forward=False ) bwd_micro_batch_id = self._record_bwd_micro_step( bwd_virtual_pp_rank ) bwd_job = core.Job( self.BACKWARD + "_b" + str(bwd_virtual_pp_rank) ) bwd_job.set_micro_batch_id(bwd_micro_batch_id) job_list.append(bwd_job) else: for micro_step in range(steady_steps): fwd_micro_step = micro_step + warmup_steps fwd_virtual_pp_rank = _get_virtual_pp_rank( fwd_micro_step, forward=True ) fwd_micro_batch_id = self._record_fwd_micro_step( fwd_virtual_pp_rank ) bw_micro_step = micro_step bwd_virtual_pp_rank = _get_virtual_pp_rank( bw_micro_step, forward=False ) bwd_micro_batch_id = self._record_bwd_micro_step( bwd_virtual_pp_rank ) for job_type in self.jobs_in_stable_phase_in_pir: if job_type.startswith(self.FORWARD) or job_type.startswith( self.RECV_FORWARD ): job = core.Job(job_type + str(fwd_virtual_pp_rank)) job.set_micro_batch_id(fwd_micro_batch_id) else: job = core.Job(job_type + str(bwd_virtual_pp_rank)) job.set_micro_batch_id(bwd_micro_batch_id) job_list.append(job) for micro_step in range(steady_steps, total_num_steps): virtual_pp_rank = _get_virtual_pp_rank(micro_step, forward=False) micro_batch_id = self._record_bwd_micro_step(virtual_pp_rank) if real_split_backward: bwd_job = core.Job(self.BACKWARD + "_b" + str(virtual_pp_rank)) bwd_job.set_micro_batch_id(micro_batch_id) job_list.append(bwd_job) else: bwd_job = core.Job(self.BACKWARD + str(virtual_pp_rank)) send_bwd_job = core.Job( self.SEND_BACKWARD + str(virtual_pp_rank) ) bwd_job.set_micro_batch_id(micro_batch_id) send_bwd_job.set_micro_batch_id(micro_batch_id) job_list.append(bwd_job) job_list.append(send_bwd_job) # TODO(lizhiyu): Inserting 'backward_b' and 'backward_w' interleavedly can decrease the memory, # but it reduces the speed. We should find the better way to use the code here. # next_virtual_pp_rank = _get_virtual_pp_rank(micro_step + 1, forward=False) # if next_virtual_pp_rank != virtual_pp_rank: # for micro_batch_id in range(0, accumulate_steps): # w_job = core.Job(BACKWARD + "_w" + str(virtual_pp_rank)) # w_job.set_micro_batch_id(micro_batch_id) # job_list.append(w_job) if real_split_backward: for chunk_id in range(num_model_chunks - 1, -1, -1): for micro_batch_id in range(0, accumulate_steps): if ( self._real_overlap_sharding_reduce and micro_batch_id == accumulate_steps - 1 ): w_job = core.Job( self.BACKWARD + "_w" + str(chunk_id) + self.reduce_comm_suffix ) else: w_job = core.Job(self.BACKWARD + "_w" + str(chunk_id)) w_job.set_micro_batch_id(micro_batch_id) job_list.append(w_job) job_types = [job.type() for job in job_list] logger.debug(f"The VPP job list: {job_types}") opt_job = core.Job(self.OPT) job_list.append(opt_job) return job_list def _pir_split_matmul_grad_ops_to_matmul(self, program): for block in program.blocks: matmul_grad_op_idx = [] ops = block.ops for i, op_i in enumerate(ops): if ( op_i.name() == "pd_op.matmul_grad" and not op_i.has_attr("trans_x") and not op_i.has_attr("trans_y") ): matmul_grad_op_idx.append(i) for matmul_grad_id in reversed(matmul_grad_op_idx): _pir_split_matmul_grad_to_matmul(block, matmul_grad_id) def _partial_programs(self, program): raise RuntimeError("Not support old IR for VPP") def _partial_pir_programs(self, program): num_model_chunks = self.get_attr("vpp_degree") enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap") split_backward = self.get_attr("split_backward", False) accumulate_steps = self.get_attr("num_micro_batches") num_stages = self.get_attr("pp_degree") if accumulate_steps != num_stages: split_backward = False assert not enable_send_recv_overlap, ( "PIR does not support VPP with enable_send_recv_overlap yet." ) if split_backward: self._pir_split_matmul_grad_ops_to_matmul(program) types, sub_program_list = self._pir_program_for_vpp( program, num_model_chunks, split_backward, enable_send_recv_overlap ) for i in range(len(types)): logger.debug( f"type = {types[i]}, sub_programs = {sub_program_list[i]}\n" ) return types, sub_program_list def _pir_program_for_vpp( self, program, num_model_chunks, split_bw=False, enable_send_recv_overlap=False, ): _pir_overlap_send_recv(program) oprole_names = [ "recv_forward", "forward", "backward", "send_backward", "optimizer", ] if split_bw: oprole_names = ["forward", "backward_b", "backward_w", "optimizer"] program_types, programs = self._split_program_for_vpp( program, num_model_chunks, oprole_names, split_bw=split_bw ) return program_types, programs def _split_program_for_vpp( self, program, num_model_chunks, oprole_names, split_bw=False ): place = _get_device() if isinstance(place, paddle.framework.CUDAPlace): place = paddle.framework.CUDAPlace( paddle.distributed.ParallelEnv().dev_id ) cur_place = paddle.base.libpaddle.Place() cur_place.set_place(place) def get_var_name(op_idx, result_idx): result_value = all_ops[op_idx].result(result_idx) all_used_ops = result_value.all_used_ops() shadow_output_op_used = None for op in all_used_ops: if op.name() == "builtin.shadow_output": shadow_output_op_used = op if shadow_output_op_used is not None: var_name = shadow_output_op_used.attrs()["output_name"] else: var_name = f"var_{op_idx}_{all_ops[op_idx].name()}_{result_idx}" return var_name def add_persistable_var(op_idx, program_type): all_program_types = list(type_to_program.keys()) following_program_types = all_program_types[ all_program_types.index(program_type) + 1 : ] op_num_results = type_to_ops[program_type][op_idx].num_results() op_name = type_to_ops[program_type][op_idx].name() for idx in range(op_num_results): var_name = None for type in reversed(following_program_types): op_result = type_to_ops[type][op_idx].result(idx) if op_result.use_empty(): continue # if this op's output is used, create the persistable # var to be used in other programs. if var_name is None: if op_name in ["pd_op.data", "builtin.parameter"]: var_name = op_result.name else: var_name = get_var_name(op_idx, idx) if "var_" in var_name: paddle.pir.set_insertion_point_after( type_to_ops[program_type][op_idx] ) paddle._C_ops.set_persistable_value( type_to_ops[program_type][op_idx].result( idx ), var_name, ) self._add_dependency_if_necessary( type_to_ops, program_type, type, op_idx, idx, var_name ) program_block = type_to_program[type].global_block() new_result_var = program_block.add_kwarg( var_name, op_result.type() ) new_result_var.place_attr = cur_place new_result_var.persistable = op_result.persistable type_to_ops[type][op_idx].result(idx).replace_all_uses_with( new_result_var ) for type in following_program_types: type_to_ops[type][op_idx].erase() type_to_program = OrderedDict() type_to_ops = OrderedDict() # Step1: create programs and ops for each type if not split_bw: chunk_ids = list(range(num_model_chunks)) # Forward process for chunk_id in chunk_ids: for job_type in ["recv_forward", "forward"]: name, prog, ops = _create_program_and_ops( program, job_type, chunk_id ) type_to_program[name] = prog type_to_ops[name] = ops # Backward process for chunk_id in reversed(chunk_ids): for job_type in ["backward", "send_backward"]: name, prog, ops = _create_program_and_ops( program, job_type, chunk_id ) type_to_program[name] = prog type_to_ops[name] = ops # Optimizer name, prog, ops = _create_program_and_ops(program, "optimizer") type_to_program[name] = prog type_to_ops[name] = ops else: for type in oprole_names: if type == "optimizer": type_to_program["optimizer"] = program.clone() type_to_ops["optimizer"] = ( type_to_program["optimizer"].global_block().ops ) else: chunk_ids = list(range(num_model_chunks)) if "backward" in type: chunk_ids.reverse() for chunk_id in chunk_ids: type_to_program[type + str(chunk_id)] = program.clone() type_to_ops[type + str(chunk_id)] = ( type_to_program[type + str(chunk_id)] .global_block() .ops ) # Step2: delete the ops not belong to the type # 1. delete ops # 2. add persistable var used between multiple programs all_ops = program.global_block().ops chunk_ids = list(range(num_model_chunks)) bwd_pattern_ops_type = [] for idx in range(len(all_ops) - 1, -1, -1): op = all_ops[idx] op_role = op.op_role op_chunk_id = op.chunk_id # Step2.1: infer chunk_id for ops that don't have chunk_id if op_role != int(OpRole.Optimize) and op_chunk_id == -1: op_chunk_id = infer_chunk_id(idx, all_ops, False) if op_chunk_id == -1: raise ValueError( f"Cannot infer chunk_id for op {op.name()} at index {idx}" ) # Step2.2: identify the job_type of the op if op_role == int(OpRole.Optimize): job_type = "optimizer" elif op_role == int(OpRole.Backward) and split_bw: if len(bwd_pattern_ops_type) == 0: bwd_pattern_ops_type = _pir_get_backward_op_type( all_ops, idx ) job_type = bwd_pattern_ops_type.pop() elif op_role == int(OpRole.Backward) and (not split_bw): if op.name() == "pd_op.send_v2": job_type = "send_backward" else: job_type = "backward" elif op_role == int(OpRole.Forward): if op.name() == "pd_op.recv_v2" and (not split_bw): job_type = "recv_forward" else: job_type = "forward" else: raise ValueError( f"The op[{op.name()}]'s op role: {op_role} isn't one of recv_forward, forward, backward, send_backward or Optimizer." ) # Step2.3: delete ops not belong to the type if not split_bw: current_type = ( job_type if job_type == "optimizer" else job_type + str(op_chunk_id) ) # Get the position of the current type in type_to_program all_types = list(type_to_ops.keys()) current_idx = all_types.index(current_type) # Delete all ops before the current type for type_name in all_types[:current_idx]: type_to_ops[type_name][idx].erase() else: for type in oprole_names: if type == job_type: break if type != "optimizer": for chunk_id in chunk_ids: type_to_ops[type + str(chunk_id)][idx].erase() else: type_to_ops[type][idx].erase() chunk_order = range(0, op_chunk_id) if "backward" in job_type: chunk_order = range(num_model_chunks - 1, op_chunk_id, -1) for chunk_id in chunk_order: type_to_ops[job_type + str(chunk_id)][idx].erase() # Step2.4: add persistable var used between multiple programs if job_type != "optimizer": add_persistable_var(idx, job_type + str(op_chunk_id)) return list(type_to_program.keys()), list(type_to_program.values())