wehub-resource-sync
604 lines
25 KiB
Python
604 lines
25 KiB
Python
# 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())
|