wehub-resource-sync
1922 lines
74 KiB
Python
1922 lines
74 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 collections
|
|
import logging
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle import pir
|
|
from paddle.autograd.backward_utils import ValueDict
|
|
from paddle.base.framework import EagerParamBase, pir_op_role_guard
|
|
from paddle.base.log_helper import get_logger
|
|
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
|
from paddle.distributed.passes.pass_base import PassContext, new_pass
|
|
from paddle.distributed.passes.pass_utils import infer_chunk_id
|
|
|
|
from .mix_to_dist_pass import dist_skip_op_list
|
|
from .process_group import get_process_group
|
|
from .reshard_funcs.base_reshard_func import (
|
|
choose_reshard_func,
|
|
copy_dist_attr_with_new_member,
|
|
copy_op_attr_with_new_member,
|
|
copy_process_mesh_with_new_member,
|
|
)
|
|
from .reshard_funcs.reshard_func_register import register_reshard_funcs
|
|
from .utils import (
|
|
_complete_op_dist_attr,
|
|
fuse_param_func,
|
|
get_pp_stage_by_pp_degree,
|
|
get_pp_stage_by_process_mesh,
|
|
get_sub_process_mesh_by_program,
|
|
partition_skip_op_list,
|
|
update_pylayer_output,
|
|
)
|
|
|
|
_logger = get_logger(
|
|
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
|
)
|
|
|
|
register_reshard_funcs()
|
|
|
|
amp_ops = ["pd_op.check_finite_and_unscale_", "pd_op.update_loss_scaling_"]
|
|
|
|
|
|
def reshard_single_value(program, op, operand, attr):
|
|
prev_var = operand.source()
|
|
if prev_var.is_dist() and prev_var.dist_attr() != attr:
|
|
operand_attr = attr.as_tensor_dist_attr()
|
|
paddle.pir.set_insertion_point(op)
|
|
with pir_op_role_guard(op.op_role):
|
|
# fold reshard
|
|
if prev_var.get_defining_op().name() == 'dist_op.reshard':
|
|
prev_reshard = prev_var.get_defining_op()
|
|
prev_reshard_input = prev_reshard.operand_source(0)
|
|
prev_reshard_result = prev_reshard.result(0)
|
|
# skil global to sub mesh reshard
|
|
if (
|
|
prev_reshard_input.dist_attr().process_mesh.ndim
|
|
== prev_reshard_result.dist_attr().process_mesh.ndim
|
|
):
|
|
if prev_reshard_input.dist_attr() == operand_attr:
|
|
return prev_reshard_input
|
|
reshard_var = paddle._C_ops.reshard_v2(
|
|
prev_reshard_input, operand_attr
|
|
)
|
|
return reshard_var
|
|
# insert reshard
|
|
reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
|
|
return reshard_var
|
|
return prev_var
|
|
|
|
|
|
def reshard_combine_value(program, op, operand, attr):
|
|
prev_var = operand.source()
|
|
|
|
assert prev_var.get_defining_op().name() == 'builtin.combine', (
|
|
f"TensorList must be defined by builtin.combine op, but is {prev_var.get_defining_op().name()}."
|
|
)
|
|
|
|
combine_op = prev_var.get_defining_op()
|
|
array_attr = attr.as_array_attr()
|
|
|
|
assert len(combine_op.operands()) == len(array_attr), (
|
|
"The number of combine op operands and the number of dist array_attr are not equal in op"
|
|
)
|
|
|
|
reshard_vars = []
|
|
for inner_operand, inner_attr in zip(combine_op.operands(), array_attr):
|
|
reshard_vars.append(
|
|
reshard_single_value(program, op, inner_operand, inner_attr)
|
|
)
|
|
|
|
paddle.pir.set_insertion_point(op)
|
|
with pir_op_role_guard(op.op_role):
|
|
combine_value = paddle._C_ops.builtin_combine(reshard_vars)
|
|
return combine_value
|
|
|
|
|
|
def apply_partition_pass(program, block=None):
|
|
if block is None:
|
|
block = program.global_block()
|
|
for op in block.ops:
|
|
for sub_block in op.blocks():
|
|
apply_partition_pass(program, block=sub_block)
|
|
|
|
if op.dist_attr is None:
|
|
continue
|
|
if op.name() in partition_skip_op_list:
|
|
continue
|
|
|
|
assert len(op.operands()) == len(op.dist_attr.operands()), (
|
|
f"The number of operands and the number of op_dist_attr's operands are not equal in op: {op}"
|
|
)
|
|
assert len(op.results()) == len(op.dist_attr.results()), (
|
|
f"The number of results and the number of op_dist_attr's results are not equal in op: {op}"
|
|
)
|
|
|
|
# deal with inplace value
|
|
for out_idx, in_idx in paddle.core.pir.get_op_inplace_info(op).items():
|
|
ref_op_role = op.op_role
|
|
|
|
operand = op.operand(in_idx)
|
|
operand_attr = op.dist_attr.operand(in_idx)
|
|
prev_var = operand.source()
|
|
if (
|
|
not prev_var.is_dist()
|
|
or operand_attr == prev_var.dist_attr()
|
|
or not prev_var.persistable
|
|
):
|
|
continue
|
|
|
|
assert not prev_var.is_combine(), (
|
|
f"The current partition pass not support inplace value of {op} is tensor list."
|
|
)
|
|
|
|
operand_attr = operand_attr.as_tensor_dist_attr()
|
|
|
|
# reshard input
|
|
paddle.pir.set_insertion_point(op)
|
|
with pir_op_role_guard(ref_op_role):
|
|
reshard_var = paddle._C_ops.reshard_v2(prev_var, operand_attr)
|
|
operand.set_source(reshard_var)
|
|
|
|
result = op.result(out_idx)
|
|
result_attr = op.dist_attr.result(out_idx).as_tensor_dist_attr()
|
|
assert operand_attr == result_attr, (
|
|
f"For inplace value, The operend dist attr should be equal to result dist attr , please check your infer_spmd func of {op}"
|
|
)
|
|
|
|
# reshard output
|
|
paddle.pir.set_insertion_point_after(op)
|
|
old_dist_attr = result.dist_attr()
|
|
result.update_dist_attr(result_attr)
|
|
|
|
with pir_op_role_guard(ref_op_role):
|
|
prev_op = prev_var.get_defining_op()
|
|
|
|
# reshard output to assign out input
|
|
reshard_var_1 = paddle._C_ops.reshard_v2(
|
|
result, prev_var.dist_attr()
|
|
)
|
|
assign_out = paddle._C_ops.assign_out_(reshard_var_1, prev_var)
|
|
assign_out.get_defining_op().dist_attr = (
|
|
copy_op_attr_with_new_member(
|
|
assign_out.get_defining_op().dist_attr,
|
|
new_chunk_id=op.dist_attr.chunk_id,
|
|
)
|
|
)
|
|
|
|
if old_dist_attr == result.dist_attr():
|
|
continue
|
|
|
|
reshard_var_2 = reshard_var_1
|
|
if old_dist_attr != reshard_var_1.dist_attr():
|
|
with pir_op_role_guard(ref_op_role):
|
|
reshard_var_2 = paddle._C_ops.reshard_v2(
|
|
result, old_dist_attr
|
|
)
|
|
|
|
result.replace_all_uses_with(reshard_var_1)
|
|
reshard_var_1.get_defining_op().operand(0).set_source(result)
|
|
reshard_var_2.get_defining_op().operand(0).set_source(result)
|
|
|
|
for operand, attr in zip(op.operands(), op.dist_attr.operands()):
|
|
if not attr:
|
|
continue
|
|
prev_var = operand.source()
|
|
if prev_var.is_combine():
|
|
operand.set_source(
|
|
reshard_combine_value(program, op, operand, attr)
|
|
)
|
|
else:
|
|
operand.set_source(
|
|
reshard_single_value(program, op, operand, attr)
|
|
)
|
|
prev_op = prev_var.get_defining_op()
|
|
if prev_op and prev_op.num_results() == 1 and prev_var.use_empty():
|
|
prev_op.erase()
|
|
|
|
for var, attr in zip(op.results(), op.dist_attr.results()):
|
|
if (
|
|
attr
|
|
and var.initialized()
|
|
and var.is_dist()
|
|
and var.dist_attr() != attr
|
|
):
|
|
paddle.pir.set_insertion_point_after(op)
|
|
old_dist_attr = var.dist_attr()
|
|
var.update_dist_attr(attr.as_tensor_dist_attr())
|
|
|
|
# insert reshard
|
|
with pir_op_role_guard(op.op_role):
|
|
reshard_var = paddle._C_ops.reshard_v2(var, old_dist_attr)
|
|
var.replace_all_uses_with(reshard_var)
|
|
reshard_var.get_defining_op().operand(0).set_source(var)
|
|
var.get_defining_op().set_bool_attr(
|
|
"replace_all_uses_with_reshard_var", True
|
|
)
|
|
|
|
|
|
class ReshardPasses:
|
|
@staticmethod
|
|
def decompose_reshard_pass(dist_program):
|
|
# split composed reshard op into atomic reshard ops, which would increase the opportunity of reshard Re-Use in following fold_reshard_pass.
|
|
del_ops = []
|
|
for op in dist_program.global_block().ops:
|
|
if op.name() != 'dist_op.reshard':
|
|
continue
|
|
input = op.operand_source(0)
|
|
result = op.result(0)
|
|
|
|
# split the reshard compose p2p and collective into one p2p reshard and one collective reshard.
|
|
# avoid global to sub mesh case
|
|
if (
|
|
(
|
|
input.dist_attr().process_mesh
|
|
!= result.dist_attr().process_mesh
|
|
)
|
|
and input.dist_attr().process_mesh.ndim
|
|
== result.dist_attr().process_mesh.ndim
|
|
):
|
|
if (
|
|
input.dist_attr().placements
|
|
!= result.dist_attr().placements
|
|
):
|
|
ref_op_role = op.op_role
|
|
with pir_op_role_guard(ref_op_role):
|
|
intermediate_dist_attr = copy_dist_attr_with_new_member(
|
|
input.dist_attr(),
|
|
new_process_mesh=result.dist_attr().process_mesh,
|
|
)
|
|
intermediate_dist_type = (
|
|
paddle.base.libpaddle.pir.cvt_to_dist_type(
|
|
input.type(), intermediate_dist_attr
|
|
)
|
|
)
|
|
paddle.pir.set_insertion_point(op)
|
|
intermediate_var = paddle._C_ops.reshard_v2(
|
|
input, intermediate_dist_attr
|
|
)
|
|
new_reshard_result = paddle._C_ops.reshard_v2(
|
|
intermediate_var, result.dist_attr()
|
|
)
|
|
result.replace_all_uses_with(new_reshard_result)
|
|
del_ops.append(op)
|
|
|
|
for op in del_ops:
|
|
_logger.info(f"[Reshard Pass] atomic composed reshard op: {op!s}")
|
|
op.erase()
|
|
|
|
@staticmethod
|
|
def fold_reshard_pass(dist_program):
|
|
del_ops = []
|
|
value_dict = ValueDict()
|
|
for op in dist_program.global_block().ops:
|
|
if op.name() != 'dist_op.reshard':
|
|
continue
|
|
input = op.operand_source(0)
|
|
result = op.result(0)
|
|
if input.type() == result.type():
|
|
result.replace_all_uses_with(input)
|
|
del_ops.append(op)
|
|
continue
|
|
if input not in value_dict:
|
|
value_dict[input] = [(result.type(), result)]
|
|
continue
|
|
no_find = True
|
|
for type, val in value_dict[input]:
|
|
if type == result.type():
|
|
result.replace_all_uses_with(val)
|
|
del_ops.append(op)
|
|
no_find = False
|
|
break
|
|
if no_find:
|
|
value_dict[input].append((result.type(), result))
|
|
for op in del_ops:
|
|
op.erase()
|
|
|
|
@staticmethod
|
|
def reshard_op_pass(dist_program, global_params_grads=None, block=None):
|
|
if block is None:
|
|
block = dist_program.global_block()
|
|
for op in block.ops:
|
|
for sub_block in op.blocks():
|
|
ReshardPasses.reshard_op_pass(dist_program, block=sub_block)
|
|
if op.name() == 'dist_op.reshard':
|
|
var = op.operand_source(0)
|
|
|
|
op_dist_attr = op.dist_attr
|
|
src_dist_attr = op_dist_attr.operand(0).as_tensor_dist_attr()
|
|
dst_dist_attr = op_dist_attr.result(0).as_tensor_dist_attr()
|
|
|
|
assert (
|
|
not var.initialized() or var.dist_attr() == src_dist_attr
|
|
), (
|
|
f"The dist_attr of reshard op's input and operand should be equal, but got {var.dist_attr()} and {src_dist_attr}"
|
|
)
|
|
|
|
if src_dist_attr == dst_dist_attr:
|
|
op.result(0).replace_all_uses_with(var)
|
|
if global_params_grads is not None:
|
|
for idx, (p, g) in enumerate(global_params_grads):
|
|
if g is not None and g.is_same(op.result(0)):
|
|
global_params_grads[idx] = (p, var)
|
|
op.erase()
|
|
continue
|
|
|
|
paddle.pir.set_insertion_point(op)
|
|
ref_op_role = op.op_role
|
|
|
|
all_to_all_dim = (
|
|
dist.auto_parallel.moe_utils._specific_alltoall_dim(
|
|
var,
|
|
dst_dist_attr.process_mesh,
|
|
dst_dist_attr.placements_attr,
|
|
)
|
|
)
|
|
|
|
if all_to_all_dim is not None:
|
|
with pir_op_role_guard(ref_op_role):
|
|
out_value = (
|
|
dist.auto_parallel.moe_utils._pir_nd_mesh_all2all(
|
|
op.operand_source(0),
|
|
op.result(0).type(),
|
|
dst_dist_attr.process_mesh,
|
|
dst_dist_attr.placements_attr,
|
|
all_to_all_dim,
|
|
)
|
|
)
|
|
else:
|
|
reshard_func = choose_reshard_func(
|
|
src_dist_attr, dst_dist_attr
|
|
)
|
|
assert reshard_func is not None, (
|
|
f'There is no reshard function that matches src_dist_attr: {src_dist_attr} and dst_dist_attr: {dst_dist_attr}, {var.get_defining_op()}'
|
|
)
|
|
|
|
with pir_op_role_guard(ref_op_role):
|
|
out_value = reshard_func.reshard(
|
|
src_dist_attr,
|
|
dst_dist_attr,
|
|
op.operand_source(0),
|
|
op.result(0).type(),
|
|
)
|
|
|
|
if out_value is not None:
|
|
op.result(0).replace_all_uses_with(out_value)
|
|
|
|
if op.result(0).use_empty():
|
|
if global_params_grads is not None:
|
|
for idx, (p, g) in enumerate(global_params_grads):
|
|
if g is not None and g.is_same(op.result(0)):
|
|
global_params_grads[idx] = (
|
|
(p, out_value)
|
|
if out_value is not None
|
|
else (p, var)
|
|
)
|
|
op.erase()
|
|
|
|
@staticmethod
|
|
def apply_reshard_pass(dist_program, global_params_grads=None):
|
|
ReshardPasses.decompose_reshard_pass(dist_program)
|
|
ReshardPasses.fold_reshard_pass(dist_program)
|
|
ReshardPasses.reshard_op_pass(dist_program, global_params_grads)
|
|
|
|
|
|
# Replace the specific MoE-related dist op with the
|
|
# executable op in the dense program. In expert parallelism
|
|
# of the MoE model, the process mesh of each expert is
|
|
# different. Two specific apis are used to transform the
|
|
# input tensor's global process mesh to the experts' local
|
|
# process meshes, which will add two dist ops in the program.
|
|
# The following two functions are used to replace the two
|
|
# dist ops with the executable share_data_ ops.
|
|
def replace_moe_sub_mesh_tensors(op):
|
|
cur_rank = paddle.distributed.get_rank()
|
|
in_value = op.operand_source(0)
|
|
out_value = None
|
|
out_idx = -1
|
|
for idx, val in enumerate(op.results()):
|
|
val_mesh = val.dist_attr().process_mesh
|
|
if cur_rank in val_mesh.process_ids:
|
|
assert out_value is None, (
|
|
f'{op} has more than one results on rank {cur_rank}'
|
|
)
|
|
out_value = val
|
|
out_idx = idx
|
|
|
|
paddle.pir.set_insertion_point(op)
|
|
local_value = paddle._C_ops.share_data_(in_value)
|
|
local_value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
|
out_value.type(), out_value.dist_attr()
|
|
)
|
|
local_value.set_type(local_value_type)
|
|
out_value.replace_all_uses_with(local_value)
|
|
|
|
op_dist_attr = op.dist_attr
|
|
share_data_op = local_value.get_defining_op()
|
|
share_data_op.dist_attr = (
|
|
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
|
op_dist_attr.process_mesh,
|
|
[op_dist_attr.operand(0).as_tensor_dist_attr()],
|
|
[op_dist_attr.result(out_idx).as_tensor_dist_attr()],
|
|
)
|
|
)
|
|
for val in op.results():
|
|
if not val.use_empty():
|
|
update_pylayer_output(val)
|
|
assert all(val.use_empty() for val in op.results())
|
|
op.erase()
|
|
|
|
|
|
def remove_sub_block_unused_inputs(op):
|
|
inputs_size = op.operand_source.num_operands()
|
|
inputs = [op.operand_source(i) for i in range(inputs_size)]
|
|
# remove unused inputs
|
|
|
|
|
|
class RemovePasses:
|
|
@staticmethod
|
|
def remove_other_rank_op_pass(dist_program):
|
|
# pruning op and value not belong to cur rank
|
|
def prune_op(block):
|
|
cur_rank = paddle.distributed.get_rank()
|
|
reverse_block_ops = block.ops[::-1]
|
|
skip_idx = 0
|
|
for idx, op in enumerate(reverse_block_ops):
|
|
if idx < skip_idx:
|
|
continue
|
|
skip_idx += 1
|
|
|
|
if op.name() == "dist_op.moe_sub_mesh_tensors":
|
|
replace_moe_sub_mesh_tensors(op)
|
|
continue
|
|
elif op.name() == "dist_op.moe_global_mesh_tensor":
|
|
replace_moe_global_mesh_tensor(op)
|
|
continue
|
|
elif op.name() == "cf.tuple_push":
|
|
stack_create_op = op.operand_source(0).get_defining_op()
|
|
if stack_create_op.result(2).use_empty():
|
|
op.erase()
|
|
continue
|
|
elif op.name() == "cf.yield":
|
|
continue
|
|
elif op.name() == "pd_op.pylayer":
|
|
# if the pylayer op is not on the current rank, we should delete it
|
|
is_cur_rank = False
|
|
for pylayer_block in list(op.blocks())[::-1]:
|
|
for sub_block_op in pylayer_block.ops:
|
|
if (
|
|
sub_block_op.dist_attr
|
|
and cur_rank
|
|
in sub_block_op.dist_attr.process_mesh.process_ids
|
|
):
|
|
is_cur_rank = True
|
|
break
|
|
if not is_cur_rank:
|
|
op.erase()
|
|
continue
|
|
for pylayer_block in list(op.blocks())[::-1]:
|
|
prune_op(pylayer_block)
|
|
# update pylayer op's inputs
|
|
op.as_pylayer_op().update_input()
|
|
continue
|
|
elif op.name() == "dist_op.dtensor_from_local":
|
|
dtensor_to_local_idx = idx
|
|
for i in range(idx, len(reverse_block_ops)):
|
|
if (
|
|
reverse_block_ops[i].name()
|
|
== "dist_op.dtensor_to_local"
|
|
):
|
|
dtensor_to_local_idx = i
|
|
break
|
|
if (
|
|
op.dist_attr
|
|
and cur_rank
|
|
not in op.dist_attr.process_mesh.process_ids
|
|
):
|
|
for i in range(idx, dtensor_to_local_idx + 1):
|
|
reverse_block_ops[i].erase()
|
|
skip_idx = dtensor_to_local_idx + 1
|
|
continue
|
|
elif op.name() in partition_skip_op_list:
|
|
can_delete = True
|
|
for val in op.results():
|
|
if not val.use_empty():
|
|
can_delete = False
|
|
if can_delete:
|
|
op.erase()
|
|
continue
|
|
|
|
if (
|
|
op.dist_attr
|
|
and cur_rank not in op.dist_attr.process_mesh.process_ids
|
|
):
|
|
op.erase()
|
|
elif op.name() == "dist_op.reshard":
|
|
assert op.result(0).use_empty(), (
|
|
f'There should not have useful dist.reshard op in remove_other_rank_op_pass. but find : {op}'
|
|
)
|
|
op.erase()
|
|
|
|
prune_op(dist_program.global_block())
|
|
|
|
# merge pd.data ops for
|
|
lr_ops = []
|
|
lr_parameters = []
|
|
for op in dist_program.global_block().ops[::-1]:
|
|
if (
|
|
op.name() == 'pd_op.data'
|
|
and "learning_rate" in op.attrs()["name"]
|
|
):
|
|
lr_ops.append(op)
|
|
if (
|
|
op.name() == 'builtin.parameter'
|
|
and "learning_rate" in op.attrs()["parameter_name"]
|
|
):
|
|
lr_parameters.append(op)
|
|
|
|
if len(lr_ops) > 1:
|
|
lr_value = lr_ops[0].result(0)
|
|
for op in lr_ops[1:]:
|
|
lr = op.result(0)
|
|
lr.replace_all_uses_with(lr_value)
|
|
op.erase()
|
|
|
|
if len(lr_parameters) > 1:
|
|
lr_value = lr_parameters[0].result(0)
|
|
for op in lr_parameters[1:]:
|
|
lr = op.result(0)
|
|
lr.replace_all_uses_with(lr_value)
|
|
op.erase()
|
|
for keyword, argument in dist_program.global_block().kwargs().items():
|
|
if argument.use_empty():
|
|
dist_program.global_block().erase_kwarg(keyword)
|
|
|
|
@staticmethod
|
|
def remove_no_need_in_startup(startup_program, main_program):
|
|
# 1. vars used in main_program
|
|
main_program_var_names = []
|
|
for key in main_program.global_block().kwargs():
|
|
main_program_var_names.append(key)
|
|
for op in main_program.global_block().ops:
|
|
for var in op.operands_source():
|
|
if var.has_name:
|
|
main_program_var_names.append(var.name)
|
|
for var in op.results():
|
|
if var.has_name:
|
|
main_program_var_names.append(var.name)
|
|
# 2. remove var op not used in main_program
|
|
for op in startup_program.global_block().ops:
|
|
for var in op.operands_source():
|
|
if var.has_name and var.name not in main_program_var_names:
|
|
op.erase()
|
|
|
|
# 3. dead code elimination
|
|
pm = pir.PassManager()
|
|
pm.add_pass('dead_code_elimination_pass', {})
|
|
pm.run(startup_program)
|
|
for op in startup_program.global_block().ops:
|
|
if op.name() == "pd_op.coalesce_tensor_":
|
|
if op.result(0).use_empty() and op.result(1).use_empty():
|
|
op.erase()
|
|
pm.run(startup_program)
|
|
|
|
@staticmethod
|
|
def remove_other_rank_input_output_pass(dist_program):
|
|
'''
|
|
Pruning value not belong to cur rank especially used for check_finite_and_unscale
|
|
and update_loss_scaling op in amp.
|
|
|
|
For example, w0 on mesh0, w1 on mesh1, before pass, the ops is:
|
|
[w0_g, w1_g], is_finite = check_finite_and_scale([w0_g, w1_g], loss_scaling)
|
|
after pass, on mesh0, the op is:
|
|
[w0_g], is_finite = check_finite_and_scale([w0_g], loss_scaling)
|
|
|
|
Note that here we do not set the op_dist_attr, since it is not used afterwards.
|
|
'''
|
|
cur_rank = paddle.distributed.get_rank()
|
|
for op in dist_program.global_block().ops[::-1]:
|
|
if op.name() not in amp_ops:
|
|
continue
|
|
new_vars = []
|
|
combine_op = op.operand_source(0).get_defining_op()
|
|
for inner_operand in (
|
|
op.operand_source(0).get_defining_op().operands()
|
|
):
|
|
if (
|
|
cur_rank
|
|
in inner_operand.source()
|
|
.dist_attr()
|
|
.process_mesh.process_ids
|
|
):
|
|
new_vars.append(inner_operand.source())
|
|
continue
|
|
result = op.operand_source(0).get_defining_op().result(0)
|
|
paddle.pir.set_insertion_point_after(combine_op)
|
|
res = paddle._C_ops.builtin_combine(new_vars)
|
|
res.get_defining_op().op_role = op.op_role
|
|
result.replace_all_uses_with(res)
|
|
combine_op.erase()
|
|
# since it is inplace op, set type of output as the same as input
|
|
op.result(0).set_type(res.type())
|
|
|
|
@staticmethod
|
|
def remove_other_rank_params_grads_pass(dist_program, dist_params_grads):
|
|
cur_rank_param = []
|
|
cur_rank = paddle.distributed.get_rank()
|
|
|
|
for op in dist_program.global_block().ops:
|
|
if op.name() == 'builtin.parameter':
|
|
if cur_rank in op.dist_attr.process_mesh.process_ids:
|
|
cur_rank_param.append(op.attrs()['parameter_name'])
|
|
|
|
need_remove_idx = []
|
|
for idx, (param, grad) in enumerate(dist_params_grads):
|
|
if grad is None:
|
|
continue
|
|
if param.name not in cur_rank_param:
|
|
need_remove_idx.append(idx)
|
|
|
|
for idx in need_remove_idx[::-1]:
|
|
dist_params_grads.pop(idx)
|
|
|
|
@staticmethod
|
|
def apply_all(
|
|
dist_main_program, dist_startup_program, dist_params_grads=[]
|
|
):
|
|
RemovePasses.remove_other_rank_input_output_pass(dist_main_program)
|
|
RemovePasses.remove_other_rank_params_grads_pass(
|
|
dist_main_program, dist_params_grads
|
|
)
|
|
_complete_op_dist_attr(dist_main_program)
|
|
RemovePasses.remove_other_rank_op_pass(dist_main_program)
|
|
RemovePasses.remove_no_need_in_startup(
|
|
dist_startup_program, dist_main_program
|
|
)
|
|
|
|
|
|
def replace_moe_global_mesh_tensor(op):
|
|
cur_rank = paddle.distributed.get_rank()
|
|
out_value = op.result(0)
|
|
in_value = None
|
|
in_idx = -1
|
|
for idx, val in enumerate(op.operands_source()):
|
|
val_mesh = val.dist_attr().process_mesh
|
|
if cur_rank not in val_mesh.process_ids:
|
|
continue
|
|
assert in_value is None, (
|
|
f'{op} has more than one inputs on rank {cur_rank}'
|
|
)
|
|
in_value = val
|
|
in_idx = idx
|
|
|
|
paddle.pir.set_insertion_point(op)
|
|
local_value = paddle._C_ops.share_data_(in_value)
|
|
# local_value = paddle.assign(in_value)
|
|
local_value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
|
out_value.type(), out_value.dist_attr()
|
|
)
|
|
local_value.set_type(local_value_type)
|
|
out_value.replace_all_uses_with(local_value)
|
|
|
|
op_dist_attr = op.dist_attr
|
|
share_data_op = local_value.get_defining_op()
|
|
share_data_op.dist_attr = (
|
|
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
|
op_dist_attr.process_mesh,
|
|
[op_dist_attr.operand(in_idx).as_tensor_dist_attr()],
|
|
[op_dist_attr.result(0).as_tensor_dist_attr()],
|
|
)
|
|
)
|
|
|
|
assert all(val.use_empty() for val in op.results())
|
|
op.erase()
|
|
|
|
|
|
# Note: this is the pass in the dense program
|
|
comm_ops = [
|
|
"pd_op.all_gather",
|
|
"pd_op.reduce_scatter",
|
|
]
|
|
|
|
|
|
def remove_unuseful_comm_op_pass(program):
|
|
for op in program.global_block().ops:
|
|
if op.name() in comm_ops or (
|
|
op.name() == "pd_op.all_reduce"
|
|
and op.int_attr("reduce_type")
|
|
in [dist.ReduceOp.SUM, dist.ReduceOp.MAX]
|
|
):
|
|
ring_id = op.int_attr("ring_id")
|
|
process_group = get_process_group(ring_id)
|
|
if process_group.nranks == 1:
|
|
op.result(0).replace_all_uses_with(op.operand_source(0))
|
|
op.erase()
|
|
if op.name() == "pd_op.share_data_":
|
|
if op.operand_source(0).has_one_use():
|
|
op.result(0).replace_all_uses_with(op.operand_source(0))
|
|
op.erase()
|
|
if (
|
|
op.name() == "pd_op.cast"
|
|
and op.result(0).dtype == op.operand_source(0).dtype
|
|
):
|
|
op.result(0).replace_all_uses_with(op.operand_source(0))
|
|
op.erase()
|
|
|
|
|
|
# In sequence_parallel, we need to transpose hidden_states
|
|
# from [bs, seq, hidden] to [seq, bs, hidden] to perform
|
|
# split and allgather at dim 0.
|
|
# The transpose may lead to about 3% performance
|
|
# in llama-70B model (tp4pp8).
|
|
# We found that, when bs=1, which is the common case in llm
|
|
# training, the transpose is equal to reshape.
|
|
# So, this pass is to haddle the specific case.
|
|
def eliminate_transpose_by_reshape(program):
|
|
for op in program.global_block().ops:
|
|
if (
|
|
op.name() == 'pd_op.transpose'
|
|
or op.name() == 'pd_op.transpose_grad'
|
|
):
|
|
var = op.operand(0).source()
|
|
rank = len(var.shape)
|
|
perm = op.attrs()['perm']
|
|
perm = [p + rank if p < 0 else p for p in perm]
|
|
# only support transpose dim 0 and dim 1
|
|
expected_perm = [1, 0] + [i + 2 for i in range(rank - 2)]
|
|
if perm == expected_perm and (
|
|
var.shape[0] == 1 or var.shape[1] == 1
|
|
):
|
|
paddle.pir.set_insertion_point(op)
|
|
transpose_var = op.result(0)
|
|
reshape_var = paddle._C_ops.reshape(var, transpose_var.shape)
|
|
reshape_var.get_defining_op().op_role = op.op_role
|
|
transpose_var.replace_all_uses_with(reshape_var)
|
|
op.erase()
|
|
return program
|
|
|
|
|
|
def complete_op_role(main_program, op_role_scope: list):
|
|
assert len(op_role_scope) == 3 and len(op_role_scope[0]) == 2, (
|
|
"op_role_scope should has the shape[3, 2]"
|
|
)
|
|
forward_op_start = op_role_scope[0][0]
|
|
forward_op_end = op_role_scope[0][1]
|
|
|
|
backward_op_start = op_role_scope[1][0]
|
|
backward_op_end = op_role_scope[1][1]
|
|
|
|
opt_op_start = op_role_scope[2][0]
|
|
opt_op_end = op_role_scope[2][1]
|
|
|
|
global_op_idx = 0
|
|
for blk in main_program.blocks:
|
|
for op in blk.ops:
|
|
if (
|
|
global_op_idx >= forward_op_start
|
|
and global_op_idx < forward_op_end
|
|
):
|
|
op.op_role = 0
|
|
elif (
|
|
global_op_idx >= backward_op_start
|
|
and global_op_idx < backward_op_end
|
|
):
|
|
op.op_role = 1
|
|
elif global_op_idx >= opt_op_start and global_op_idx < opt_op_end:
|
|
op.op_role = 2
|
|
else:
|
|
pass
|
|
global_op_idx += 1
|
|
|
|
|
|
def pipeline_pass(dense_main_program, dense_startup_program, pipeline_strategy):
|
|
"""
|
|
Pipeline schedule pass for auto parallel. Enables the pipeline parallel scheduling
|
|
strategies like FThenB, 1F1B, VPP, etc.
|
|
"""
|
|
import os
|
|
|
|
pass_name = pipeline_strategy.schedule_mode
|
|
assert pass_name in [
|
|
"FThenB",
|
|
"1F1B",
|
|
"VPP",
|
|
], (
|
|
f"pipeline scheduler only support FThenB, 1F1B and VPP now, but receive {pass_name}"
|
|
)
|
|
|
|
pass_attr = {}
|
|
pass_attr["num_micro_batches"] = pipeline_strategy.accumulate_steps
|
|
pass_attr["pp_degree"] = pipeline_strategy.pp_degree
|
|
pass_attr["pp_stage"] = get_pp_stage_by_pp_degree(
|
|
pipeline_strategy.pp_degree
|
|
)
|
|
pass_attr["vpp_degree"] = pipeline_strategy.vpp_degree
|
|
pass_attr["split_backward"] = pipeline_strategy.split_backward
|
|
|
|
if pass_name == "1F1B":
|
|
# TODO(Ruibiao): Move FLAGS_1f1b_backward_forward_overlap and
|
|
# FLAGS_mp_async_allreduce_in_backward to auto parallel Strategy
|
|
# after these two optimizations are available.
|
|
pass_attr["enable_backward_forward_overlap"] = int(
|
|
os.environ.get("FLAGS_1f1b_backward_forward_overlap", 0)
|
|
)
|
|
|
|
pipeline_pass = new_pass("pipeline_scheduler_" + pass_name, pass_attr)
|
|
pass_context = PassContext()
|
|
pipeline_pass.apply(
|
|
dense_main_program,
|
|
dense_startup_program,
|
|
pass_context,
|
|
)
|
|
plan = pass_context.get_attr("plan")
|
|
return plan
|
|
|
|
|
|
def _extract_seg_method(op, seg_method):
|
|
regex = re.compile(seg_method, re.IGNORECASE)
|
|
struct_name = (
|
|
op.attrs()["struct_name"] if op.has_attr("struct_name") else "/"
|
|
)
|
|
m = regex.search(struct_name)
|
|
if not m:
|
|
return None
|
|
return struct_name[m.start(0) :].split("/")[0]
|
|
|
|
|
|
def _get_seg_struct_names(ops, seg_method):
|
|
fwd_start_op_index = 0
|
|
for i, op in enumerate(ops):
|
|
if _extract_seg_method(op, seg_method):
|
|
fwd_start_op_index = i
|
|
break
|
|
|
|
total_op_num = len(ops)
|
|
fwd_end_op_index = total_op_num - 1
|
|
for i in reversed(range(total_op_num)):
|
|
if _extract_seg_method(ops[i], seg_method):
|
|
fwd_end_op_index = i
|
|
break
|
|
|
|
struct_names = collections.OrderedDict()
|
|
seg_op_mesh = collections.OrderedDict()
|
|
|
|
for i in range(fwd_start_op_index, fwd_end_op_index + 1):
|
|
if ops[i].dist_attr is None:
|
|
continue
|
|
|
|
struct_name = _extract_seg_method(ops[i], seg_method)
|
|
if struct_name:
|
|
struct_names[struct_name] = 1
|
|
if struct_name in seg_op_mesh:
|
|
assert (
|
|
seg_op_mesh[struct_name] == ops[i].dist_attr.process_mesh
|
|
), "The segment's ops should have same process_mesh."
|
|
|
|
seg_op_mesh[struct_name] = ops[i].dist_attr.process_mesh
|
|
else:
|
|
if ops[i].name() != "dist_op.reshard":
|
|
raise ValueError(
|
|
f"The op {ops[i].name()} without seg_method in its struct_name should only be reshard"
|
|
)
|
|
|
|
return list(struct_names.keys())
|
|
|
|
|
|
def _analyze_use_custom_mesh(ops, seg_method, pp_degree):
|
|
non_use_custom_mesh = True
|
|
seg_pp_stages = [-1]
|
|
|
|
for op in ops:
|
|
if _extract_seg_method(op, seg_method) and op.dist_attr:
|
|
op_mesh = op.dist_attr.process_mesh
|
|
pp_stage = get_pp_stage_by_process_mesh(op_mesh, pp_degree)
|
|
if pp_stage is None:
|
|
continue
|
|
|
|
if seg_pp_stages[-1] > pp_stage:
|
|
non_use_custom_mesh = False
|
|
break
|
|
seg_pp_stages.append(pp_stage)
|
|
|
|
if not non_use_custom_mesh:
|
|
_logger.info("Cannot Use Auto VPP")
|
|
else:
|
|
_logger.info("Using Auto VPP")
|
|
|
|
return non_use_custom_mesh
|
|
|
|
|
|
def _set_process_mesh_and_chunk_id(
|
|
op,
|
|
chunk_process_mesh,
|
|
chunk_id,
|
|
set_input_mesh=False,
|
|
set_output_mesh=False,
|
|
):
|
|
def set_var_origin_op_process_mesh(var_origin_op):
|
|
var_origin_op_input_attr = var_origin_op.dist_attr.operands()
|
|
var_origin_op_output_attr = var_origin_op.dist_attr.results()
|
|
var_origin_op_output_attr[0] = var_origin_op_output_attr[
|
|
0
|
|
].as_tensor_dist_attr()
|
|
var_origin_op_output_attr[0] = (
|
|
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
|
chunk_process_mesh,
|
|
var_origin_op_output_attr[0].dims_mapping,
|
|
var_origin_op_output_attr[0].partial_status,
|
|
)
|
|
)
|
|
|
|
var_origin_op.dist_attr = (
|
|
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
|
chunk_process_mesh,
|
|
var_origin_op_input_attr,
|
|
var_origin_op_output_attr,
|
|
0,
|
|
)
|
|
)
|
|
|
|
def get_var_process_mesh(var):
|
|
var_process_mesh = None
|
|
var_dist_attr = var.dist_attr()
|
|
|
|
def get_attr_mesh(var_dist_attr):
|
|
if var_dist_attr:
|
|
if var_dist_attr.as_array_attr():
|
|
var_array_attr = var_dist_attr.as_array_attr()
|
|
return var_array_attr[0].as_tensor_dist_attr().process_mesh
|
|
else:
|
|
return var_dist_attr.process_mesh
|
|
|
|
if var_dist_attr:
|
|
var_process_mesh = get_attr_mesh(var_dist_attr)
|
|
elif var.is_combine():
|
|
# NOTE(zhangwl): op var may is vec_type , need get var dist_attr one by one
|
|
var_list = var.type().as_vec_type()
|
|
var_list = var_list.as_list() if var_list is not None else var_list
|
|
var_attr_list = []
|
|
for combine_var in var_list:
|
|
var_dist_attr = combine_var.as_dist_type().dist_attr()
|
|
var_process_mesh = get_attr_mesh(var_dist_attr)
|
|
if var_process_mesh is not None:
|
|
return var_process_mesh
|
|
|
|
def get_var_attr_with_process_mesh(
|
|
var_dist_attr, var_origin_op, process_mesh
|
|
):
|
|
# Note(luchang): the var generated by builtin.combine will have multiple dist_attr
|
|
if var_dist_attr and var_dist_attr.as_array_attr():
|
|
var_array_attr = var_dist_attr.as_array_attr()
|
|
for i in range(len(var_array_attr)):
|
|
var_dist_attr = var_array_attr[i].as_tensor_dist_attr()
|
|
if op_mesh is not None:
|
|
if var_dist_attr.process_mesh == op_mesh:
|
|
var_array_attr[i] = copy_dist_attr_with_new_member(
|
|
var_dist_attr, new_process_mesh=process_mesh
|
|
)
|
|
else:
|
|
var_array_attr[i] = copy_dist_attr_with_new_member(
|
|
var_dist_attr, new_process_mesh=process_mesh
|
|
)
|
|
return var_array_attr
|
|
elif var_dist_attr:
|
|
if op_mesh is not None:
|
|
if var_dist_attr.process_mesh == op_mesh:
|
|
if var_origin_op.name() in [
|
|
"pd_op.data",
|
|
"builtin.parameter",
|
|
]:
|
|
set_var_origin_op_process_mesh(var_origin_op)
|
|
var_attr = copy_dist_attr_with_new_member(
|
|
var_dist_attr, new_process_mesh=process_mesh
|
|
)
|
|
return var_attr
|
|
else:
|
|
var_attr = copy_dist_attr_with_new_member(
|
|
var_dist_attr, new_process_mesh=process_mesh
|
|
)
|
|
return var_attr
|
|
return var_dist_attr
|
|
|
|
def set_var_process_mesh(var, process_mesh):
|
|
var_dist_attr = var.dist_attr()
|
|
var_origin_op = var.get_defining_op()
|
|
if var_dist_attr:
|
|
var_attr = get_var_attr_with_process_mesh(
|
|
var_dist_attr, var_origin_op, process_mesh
|
|
)
|
|
if var_attr is not None:
|
|
var.update_dist_attr(var_attr)
|
|
elif var.is_combine():
|
|
# NOTE(zhangwl): op var may is vec_type , need set var dist_attr one by one
|
|
var_list = var.type().as_vec_type()
|
|
var_list = var_list.as_list() if var_list is not None else var_list
|
|
var_attr_list = []
|
|
for combine_var in var_list:
|
|
var_dist_attr = combine_var.as_dist_type().dist_attr()
|
|
var_attr_list.append(
|
|
get_var_attr_with_process_mesh(
|
|
var_dist_attr, var_origin_op, process_mesh
|
|
)
|
|
)
|
|
var_array_attr = (
|
|
paddle.base.libpaddle.pir.create_array_dist_attribute(
|
|
var_attr_list
|
|
)
|
|
)
|
|
var.update_dist_attr(var_array_attr)
|
|
|
|
def set_attrs_process_mesh(attrs, process_mesh):
|
|
for idx, attr in enumerate(attrs):
|
|
if attr.as_array_attr():
|
|
array_attr = attr.as_array_attr()
|
|
new_array_attr = []
|
|
for i in range(len(array_attr)):
|
|
tensor_attr = array_attr[i].as_tensor_dist_attr()
|
|
new_array_attr.append(tensor_attr)
|
|
if tensor_attr and tensor_attr.process_mesh == op_mesh:
|
|
new_array_attr[i] = copy_dist_attr_with_new_member(
|
|
tensor_attr, new_process_mesh=process_mesh
|
|
)
|
|
attrs[idx] = (
|
|
paddle.base.libpaddle.pir.create_array_dist_attribute(
|
|
new_array_attr
|
|
)
|
|
)
|
|
else:
|
|
tensor_attr = attr.as_tensor_dist_attr()
|
|
if tensor_attr and tensor_attr.process_mesh == op_mesh:
|
|
attrs[idx] = copy_dist_attr_with_new_member(
|
|
tensor_attr, new_process_mesh=process_mesh
|
|
)
|
|
|
|
def set_process_mesh(vars, attrs, process_mesh):
|
|
if vars is not None:
|
|
for var in vars:
|
|
set_var_process_mesh(var, process_mesh)
|
|
if attrs is not None:
|
|
set_attrs_process_mesh(attrs, process_mesh)
|
|
|
|
op_input_vars = op.operands_source()
|
|
op_output_vars = op.results()
|
|
|
|
# NOTE(zhangwl):dist_skip_op do not have op_mesh
|
|
op_mesh = None
|
|
if op.name() in dist_skip_op_list:
|
|
input_var_process_mesh = None
|
|
# NOTE(zhangwl):dist_skip_op output_process_mesh must equal to input_process_mesh
|
|
for var in op_input_vars:
|
|
input_var_process_mesh = get_var_process_mesh(var)
|
|
if input_var_process_mesh is not None:
|
|
break
|
|
if input_var_process_mesh is not None:
|
|
set_process_mesh(op_output_vars, None, input_var_process_mesh)
|
|
return
|
|
|
|
op_dist_attr = op.dist_attr
|
|
op_mesh = op_dist_attr.process_mesh
|
|
op_input_attrs = op_dist_attr.operands()
|
|
op_output_attrs = op_dist_attr.results()
|
|
|
|
# if op in seq_chunk , vpp need set var and op chunk_process_mesh and chunk_id
|
|
if set_input_mesh:
|
|
set_process_mesh(op_input_vars, op_input_attrs, chunk_process_mesh)
|
|
if set_output_mesh:
|
|
set_process_mesh(op_output_vars, op_output_attrs, chunk_process_mesh)
|
|
if set_input_mesh or set_output_mesh:
|
|
op_mesh = chunk_process_mesh
|
|
|
|
op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
|
|
op_mesh,
|
|
op_input_attrs,
|
|
op_output_attrs,
|
|
chunk_id,
|
|
)
|
|
|
|
|
|
def complete_chunk_id(dist_program, startup_program, pipeline_strategy):
|
|
if not pipeline_strategy.enable:
|
|
return
|
|
|
|
sub_process_meshes = get_sub_process_mesh_by_program(dist_program)
|
|
pp_degree = pipeline_strategy.pp_degree
|
|
vpp_degree = pipeline_strategy.vpp_degree
|
|
seg_method = pipeline_strategy.vpp_seg_method
|
|
schedule_mode = pipeline_strategy.schedule_mode
|
|
num_chunks = pp_degree * vpp_degree
|
|
|
|
if pp_degree < 2 and vpp_degree > 1:
|
|
raise ValueError("VPP schedule mode only can be set in pipeline mode.")
|
|
if vpp_degree > 1 and (not seg_method or schedule_mode != "VPP"):
|
|
raise ValueError(
|
|
"Please set right schedule_mode and vpp_seg_method for VPP."
|
|
)
|
|
if vpp_degree < 2:
|
|
return
|
|
|
|
ReshardPasses.fold_reshard_pass(dist_program)
|
|
seg_struct_names = _get_seg_struct_names(
|
|
dist_program.global_block().ops, seg_method
|
|
)
|
|
ops = dist_program.global_block().ops
|
|
|
|
# Step2: analysis whether the pp_stage is non-decreasing among segments
|
|
# 1. if non_use_custom_mesh is True, the ops' process_mesh will be changed by vpp strategy
|
|
# 2. if non_use_custom_mesh is False, the ops's process_mesh will not be changed.
|
|
non_use_custom_mesh = _analyze_use_custom_mesh(ops, seg_method, pp_degree)
|
|
|
|
# Step3: Get op index boundary, pp_stage, chunk_id, struct_names of each segment
|
|
seg_pp_stages = [i % pp_degree for i in range(num_chunks)]
|
|
seg_chunk_ids = [i // pp_degree for i in range(num_chunks)]
|
|
seg_parts = [0]
|
|
last_struct_name = None
|
|
# stage_ids[i] represents the stage number assigned to the i-th layer.
|
|
stage_ids = []
|
|
|
|
for idx, op in enumerate(ops):
|
|
if len(seg_parts) == len(seg_struct_names):
|
|
break
|
|
struct_name = _extract_seg_method(op, seg_method)
|
|
if op.dist_attr is not None and last_struct_name != struct_name:
|
|
pp_stage = get_pp_stage_by_process_mesh(
|
|
op.dist_attr.process_mesh, pp_degree
|
|
)
|
|
if pp_stage is not None:
|
|
stage_ids.append(pp_stage)
|
|
last_struct_name = struct_name
|
|
if struct_name == seg_struct_names[len(seg_parts)]:
|
|
seg_parts.append(idx)
|
|
seg_parts.append(len(ops))
|
|
|
|
pp_stage_layer_nums = [0] * pp_degree
|
|
for i in stage_ids:
|
|
pp_stage_layer_nums[i] = pp_stage_layer_nums[i] + 1
|
|
assert all(value >= vpp_degree for value in pp_stage_layer_nums), (
|
|
"The number of layers on each pp_stage must not be less than the vpp_degree in the pp_stage to ensure that each chunk contains at least one layer."
|
|
)
|
|
|
|
seg_layer_num = [0] * num_chunks
|
|
for pp_stage in range(
|
|
0, pp_degree
|
|
): # Each pp_stage is assigned a number of layers based on user intent.
|
|
pp_stage_layer_num = pp_stage_layer_nums[pp_stage]
|
|
for i in range(0, pp_stage_layer_num):
|
|
# The pp_stage uses a Round robin scheduling algorithm to allocate layers one by one.
|
|
virtual_chunk_id = i % vpp_degree
|
|
real_chunk_id = (virtual_chunk_id) * pp_degree + pp_stage
|
|
seg_layer_num[real_chunk_id] = seg_layer_num[real_chunk_id] + 1
|
|
|
|
# Step4: Set the process_mesh of each op
|
|
seg_id = 0
|
|
reshard_ops = []
|
|
previous_seg_parts_end_idx = 0
|
|
for seg_id in range(num_chunks):
|
|
start_idx = seg_parts[previous_seg_parts_end_idx]
|
|
end_idx = seg_parts[previous_seg_parts_end_idx + seg_layer_num[seg_id]]
|
|
pp_stage = seg_pp_stages[seg_id]
|
|
chunk_id = seg_chunk_ids[seg_id]
|
|
struct_name = ",".join(
|
|
seg_struct_names[
|
|
previous_seg_parts_end_idx : previous_seg_parts_end_idx
|
|
+ seg_layer_num[seg_id]
|
|
]
|
|
)
|
|
previous_seg_parts_end_idx = (
|
|
previous_seg_parts_end_idx + seg_layer_num[seg_id]
|
|
)
|
|
process_mesh = sub_process_meshes[pp_stage]
|
|
|
|
_logger.info(
|
|
f"stage=[{pp_stage}], chunk_id=[{chunk_id}], layer_name=[{struct_name}]"
|
|
)
|
|
_logger.info(
|
|
f"start op: [{ops[start_idx].name()}], end op: [{ops[end_idx - 1].name()}]"
|
|
)
|
|
|
|
skip_idx = 0
|
|
for idx in range(start_idx, end_idx):
|
|
if idx < skip_idx:
|
|
continue
|
|
|
|
is_seg_op = _extract_seg_method(ops[idx], seg_method) is not None
|
|
set_mesh = non_use_custom_mesh & is_seg_op
|
|
|
|
if ops[idx].name() == "dist_op.reshard":
|
|
reshard_ops.append(ops[idx])
|
|
continue
|
|
elif ops[idx].name() == "dist_op.dtensor_to_local":
|
|
_set_process_mesh_and_chunk_id(
|
|
ops[idx],
|
|
process_mesh,
|
|
chunk_id,
|
|
set_input_mesh=set_mesh,
|
|
)
|
|
dtensor_from_local_idx = idx + 1
|
|
while (
|
|
ops[dtensor_from_local_idx].name()
|
|
!= "dist_op.dtensor_from_local"
|
|
):
|
|
dtensor_from_local_idx += 1
|
|
for local_op_idx in range(idx + 1, dtensor_from_local_idx):
|
|
ops[local_op_idx].set_int_attr("chunk_id", chunk_id)
|
|
|
|
_set_process_mesh_and_chunk_id(
|
|
ops[dtensor_from_local_idx],
|
|
process_mesh,
|
|
chunk_id,
|
|
set_output_mesh=set_mesh,
|
|
)
|
|
skip_idx = dtensor_from_local_idx + 1
|
|
continue
|
|
|
|
for sub_block in ops[idx].blocks():
|
|
# TODO(luchang): support condition block
|
|
pass
|
|
|
|
_set_process_mesh_and_chunk_id(
|
|
ops[idx],
|
|
process_mesh,
|
|
chunk_id,
|
|
set_input_mesh=set_mesh,
|
|
set_output_mesh=set_mesh,
|
|
)
|
|
skip_idx = idx + 1
|
|
|
|
# Step5: set right process_mesh for reshard op
|
|
for op in reshard_ops:
|
|
var = op.operand_source(0)
|
|
|
|
op_dist_attr = op.dist_attr
|
|
src_dist_attr = op_dist_attr.operand(0).as_tensor_dist_attr()
|
|
dst_dist_attr = op_dist_attr.result(0).as_tensor_dist_attr()
|
|
|
|
if src_dist_attr == dst_dist_attr:
|
|
op.result(0).replace_all_uses_with(var)
|
|
op.erase()
|
|
continue
|
|
|
|
reshard_func = choose_reshard_func(src_dist_attr, dst_dist_attr)
|
|
reshard_func_name = reshard_func.__class__.__name__
|
|
|
|
if reshard_func_name == "NdMeshReshardFunction":
|
|
new_process_mesh = var.dist_attr().process_mesh
|
|
new_src_dist_attr = copy_dist_attr_with_new_member(
|
|
src_dist_attr, new_process_mesh=new_process_mesh
|
|
)
|
|
new_dst_dist_attr = copy_dist_attr_with_new_member(
|
|
dst_dist_attr, new_process_mesh=new_process_mesh
|
|
)
|
|
op.dist_attr = copy_op_attr_with_new_member(
|
|
op_dist_attr,
|
|
new_operands=[new_src_dist_attr],
|
|
new_results=[new_dst_dist_attr],
|
|
new_process_mesh=new_process_mesh,
|
|
)
|
|
elif reshard_func_name == "GlobalToSubMeshFunction":
|
|
result_var = op.result(0)
|
|
new_process_mesh = result_var.dist_attr().process_mesh
|
|
new_dst_dist_attr = copy_dist_attr_with_new_member(
|
|
dst_dist_attr, new_process_mesh=new_process_mesh
|
|
)
|
|
op.dist_attr = copy_op_attr_with_new_member(
|
|
op_dist_attr, new_results=[new_dst_dist_attr]
|
|
)
|
|
elif reshard_func_name == "NdMeshReshardFunctionCrossMesh":
|
|
result_var = op.result(0)
|
|
src_process_mesh = var.dist_attr().process_mesh
|
|
dst_process_mesh = result_var.dist_attr().process_mesh
|
|
new_src_dist_attr = copy_dist_attr_with_new_member(
|
|
src_dist_attr, new_process_mesh=src_process_mesh
|
|
)
|
|
new_dst_dist_attr = copy_dist_attr_with_new_member(
|
|
dst_dist_attr, new_process_mesh=dst_process_mesh
|
|
)
|
|
new_process_ids = (
|
|
src_process_mesh.process_ids + dst_process_mesh.process_ids
|
|
)
|
|
new_process_mesh = copy_process_mesh_with_new_member(
|
|
op.dist_attr.process_mesh,
|
|
new_process_ids=new_process_ids,
|
|
)
|
|
|
|
op.dist_attr = copy_op_attr_with_new_member(
|
|
op_dist_attr,
|
|
new_operands=[new_src_dist_attr],
|
|
new_results=[new_dst_dist_attr],
|
|
new_process_mesh=new_process_mesh,
|
|
)
|
|
elif reshard_func_name == "SameStatusReshardFunction":
|
|
op.result(0).replace_all_uses_with(var)
|
|
op.erase()
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported reshard function: {reshard_func_name}, reshard op's dist_attr: {op.dist_attr}"
|
|
)
|
|
|
|
# Step6: add reshard op between pipeline chunks
|
|
apply_partition_pass(dist_program)
|
|
|
|
for op in startup_program.global_block().ops:
|
|
if op.name() == "builtin.set_parameter":
|
|
param_name = op.str_attr("parameter_name")
|
|
startup_param = op.operand_source(0)
|
|
param = dist_program.get_parameter_value_by_name(param_name)
|
|
if param.dist_attr():
|
|
startup_param.update_dist_attr(param.dist_attr())
|
|
|
|
ReshardPasses.fold_reshard_pass(dist_program)
|
|
|
|
|
|
def check_chunk_id(dist_program):
|
|
all_ops = dist_program.global_block().ops
|
|
for idx, op in enumerate(all_ops):
|
|
if op.op_role in [int(OpRole.Forward), int(OpRole.Backward)]:
|
|
if op.name() in dist_skip_op_list:
|
|
continue
|
|
|
|
if op.has_attr("chunk_id"):
|
|
# op between dtensor_from_local and dtensor_to_local will
|
|
# be assigned a chunk_id attribute.
|
|
continue
|
|
elif op.name() in [
|
|
"dist_op.dtensor_from_local",
|
|
"dist_op.dtensor_to_local",
|
|
]:
|
|
# dtensor_from_local and dtensor_to_local ops will be removed after
|
|
# converting the program to a dense program.
|
|
continue
|
|
elif op.dist_attr.chunk_id == -1:
|
|
if op.name() in ["pd_op.data", "builtin.parameter"]:
|
|
op.dist_attr = copy_op_attr_with_new_member(
|
|
op.dist_attr, new_chunk_id=0
|
|
)
|
|
elif op.name() in ["pd_op.full", "pd_op.full_int_array"]:
|
|
all_used_ops = op.result(0).all_used_ops()
|
|
for used_op in all_used_ops:
|
|
if used_op.dist_attr.chunk_id != -1:
|
|
op.dist_attr = copy_op_attr_with_new_member(
|
|
op.dist_attr,
|
|
new_chunk_id=used_op.dist_attr.chunk_id,
|
|
)
|
|
break
|
|
|
|
else:
|
|
op_chunk_id = infer_chunk_id(idx, all_ops)
|
|
op.dist_attr = copy_op_attr_with_new_member(
|
|
op.dist_attr, new_chunk_id=op_chunk_id
|
|
)
|
|
|
|
if op.dist_attr.chunk_id == -1:
|
|
raise ValueError(
|
|
f"The chunk_id of op[{op.name()}] is not set. Please check the chunk_id setting."
|
|
)
|
|
|
|
|
|
def check_order(op_list, order):
|
|
pointer = 0
|
|
for item in order:
|
|
if item == "pd_op.add":
|
|
while (
|
|
pointer < len(op_list)
|
|
and op_list[pointer].name() == "pd_op.add"
|
|
):
|
|
pointer += 1
|
|
else:
|
|
if pointer >= len(op_list) or op_list[pointer].name() != item:
|
|
return False
|
|
pointer += 1
|
|
return True
|
|
|
|
|
|
def is_ffn_pattern(op_list):
|
|
if len(op_list) != 3 and len(op_list) != 5:
|
|
return False
|
|
order = [
|
|
"pd_op.matmul",
|
|
"pd_op.add",
|
|
"pd_op.matmul",
|
|
"pd_op.add",
|
|
"pd_op.swiglu",
|
|
]
|
|
return check_order(op_list, order)
|
|
|
|
|
|
def is_qkv_pattern(op_list):
|
|
if len(op_list) != 9 and len(op_list) != 12:
|
|
return False
|
|
order = [
|
|
"pd_op.matmul",
|
|
"pd_op.add",
|
|
"pd_op.full_int_array",
|
|
"pd_op.reshape",
|
|
"pd_op.matmul",
|
|
"pd_op.add",
|
|
"pd_op.full_int_array",
|
|
"pd_op.reshape",
|
|
"pd_op.matmul",
|
|
"pd_op.add",
|
|
"pd_op.full_int_array",
|
|
"pd_op.reshape",
|
|
]
|
|
return check_order(op_list, order)
|
|
|
|
|
|
def get_param_op(program, param_name):
|
|
all_ops = program.global_block().ops
|
|
for i in range(len(all_ops)):
|
|
if (
|
|
all_ops[i].name() == "builtin.set_parameter"
|
|
and all_ops[i].str_attr("parameter_name") == param_name
|
|
):
|
|
return [all_ops[i], all_ops[i].operand_source(0).get_defining_op()]
|
|
|
|
|
|
@dataclass
|
|
class ParamMeta:
|
|
name: str = None
|
|
local_shape: list = None
|
|
local_num_head: int = None
|
|
local_head_dims: int = None
|
|
|
|
|
|
def fuse_attention_ffn_qkv_pass(
|
|
startup_program, main_program, concrete_program, mode="all"
|
|
):
|
|
# 0. Prepare the data structure
|
|
pir_param_names = []
|
|
dy_param_names = []
|
|
for i in range(len(concrete_program.parameters[1])):
|
|
dy_param_names.append(concrete_program.parameters[0][i].name)
|
|
pir_param_names.append(concrete_program.parameters[1][i].name)
|
|
fused_pattern_map = {"ffn": [], "qkv": []}
|
|
fusion_map = {"ffn": [], "qkv": []}
|
|
|
|
# 1. Traverse main_program, extract all ffn and qkv patterns.
|
|
all_ops = main_program.global_block().ops
|
|
for i in range(len(all_ops)):
|
|
# check ffn pattern
|
|
if mode == "all" or mode == "ffn":
|
|
pat = all_ops[i : i + 3] if i + 3 <= len(all_ops) else all_ops[i:]
|
|
if is_ffn_pattern(pat):
|
|
fused_pattern_map['ffn'].append(pat)
|
|
i = i + 3
|
|
continue
|
|
else:
|
|
pat = (
|
|
all_ops[i : i + 5] if i + 5 <= len(all_ops) else all_ops[i:]
|
|
)
|
|
if is_ffn_pattern(pat):
|
|
fused_pattern_map['ffn'].append(pat)
|
|
i = i + 5
|
|
continue
|
|
# check qkv pattern
|
|
if mode == "all" or mode == "qkv":
|
|
pat = all_ops[i : i + 9] if i + 9 <= len(all_ops) else all_ops[i:]
|
|
if is_qkv_pattern(pat):
|
|
fused_pattern_map['qkv'].append(pat)
|
|
i = i + 9
|
|
continue
|
|
else:
|
|
pat = (
|
|
all_ops[i : i + 12]
|
|
if i + 12 <= len(all_ops)
|
|
else all_ops[i:]
|
|
)
|
|
if is_qkv_pattern(pat):
|
|
fused_pattern_map['qkv'].append(pat)
|
|
i = i + 12
|
|
continue
|
|
|
|
name2pir_param_map = {}
|
|
|
|
# 2. Replace all ffn and qkv patterns with fusion patterns, and record the weights after replacement.
|
|
for pat in fused_pattern_map['ffn']:
|
|
if len(pat) == 5:
|
|
mm_gate = pat[0]
|
|
add_gate = pat[1]
|
|
mm_up = pat[2]
|
|
add_up = pat[3]
|
|
else:
|
|
mm_gate = pat[0]
|
|
add_gate = None
|
|
mm_up = pat[1]
|
|
add_up = None
|
|
|
|
fusion_w_name = f"fused_{mm_gate.operand_source(1).name}_{mm_up.operand_source(1).name}"
|
|
fusion_map["ffn"].append(
|
|
{
|
|
fusion_w_name: [
|
|
ParamMeta(mm_gate.operand_source(1).name, None, None, None),
|
|
ParamMeta(mm_up.operand_source(1).name, None, None, None),
|
|
]
|
|
}
|
|
)
|
|
|
|
fusion_w_dtype = mm_gate.operand_source(1).dtype
|
|
fusion_w_shape = mm_gate.operand_source(1).shape
|
|
fusion_w_shape[-1] += mm_up.operand_source(1).shape[-1]
|
|
fusion_w_process_mesh = mm_gate.operand_source(1).process_mesh
|
|
# Insert fusion parameter
|
|
with paddle.static.program_guard(main_program, startup_program):
|
|
fused_w = paddle.pir.core.create_parameter(
|
|
dtype=fusion_w_dtype,
|
|
shape=fusion_w_shape,
|
|
name=fusion_w_name,
|
|
process_mesh=fusion_w_process_mesh,
|
|
placements=[
|
|
paddle.distributed.Replicate(),
|
|
paddle.distributed.Shard(1),
|
|
],
|
|
initializer=paddle.nn.initializer.Constant(value=0),
|
|
)
|
|
name2pir_param_map[fusion_w_name] = fused_w
|
|
if add_gate is not None and add_up is not None:
|
|
fusion_bias_name = f"fused_{add_gate.operand_source(1).name}_{add_up.operand_source(1).name}"
|
|
fusion_map["ffn"].append(
|
|
{
|
|
fusion_bias_name: [
|
|
ParamMeta(
|
|
add_gate.operand_source(1).name, None, None, None
|
|
),
|
|
ParamMeta(
|
|
add_up.operand_source(1).name, None, None, None
|
|
),
|
|
]
|
|
}
|
|
)
|
|
|
|
fusion_bias_dtype = add_gate.operand_source(1).dtype
|
|
fusion_bias_shape = add_gate.operand_source(1).shape
|
|
fusion_bias_shape[-1] += add_up.operand_source(1).shape[-1]
|
|
fusion_bias_process_mesh = add_gate.operand_source(1).process_mesh
|
|
# Insert fusion parameter
|
|
with paddle.static.program_guard(main_program, startup_program):
|
|
fused_bias = paddle.pir.core.create_parameter(
|
|
dtype=fusion_bias_dtype,
|
|
shape=fusion_bias_shape,
|
|
name=fusion_bias_name,
|
|
process_mesh=fusion_bias_process_mesh,
|
|
placements=[
|
|
paddle.distributed.Replicate(),
|
|
paddle.distributed.Shard(0),
|
|
],
|
|
initializer=paddle.nn.initializer.Constant(value=0),
|
|
)
|
|
name2pir_param_map[fusion_bias_name] = fused_bias
|
|
|
|
# Insert dst pattern
|
|
paddle.pir.set_insertion_point_after(pat[-1])
|
|
fused_o = paddle.matmul(
|
|
mm_gate.operand_source(0),
|
|
fused_w,
|
|
transpose_x=False,
|
|
transpose_y=False,
|
|
)
|
|
fused_o.get_defining_op().copy_attrs_from(mm_gate)
|
|
if add_gate is not None and add_up is not None:
|
|
fused_o = paddle.add(fused_o, fused_bias)
|
|
fused_o.get_defining_op().copy_attrs_from(add_gate)
|
|
out = paddle.nn.functional.swiglu(fused_o)
|
|
out.get_defining_op().copy_attrs_from(pat[-1])
|
|
pat[-1].result(0).replace_all_uses_with(out)
|
|
|
|
for pat in fused_pattern_map['qkv']:
|
|
if len(pat) == 12:
|
|
mm_q = pat[0]
|
|
add_q = pat[1]
|
|
reshape_q = pat[3]
|
|
mm_k = pat[4]
|
|
add_k = pat[5]
|
|
reshape_k = pat[7]
|
|
mm_v = pat[8]
|
|
add_v = pat[9]
|
|
reshape_v = pat[11]
|
|
else:
|
|
mm_q = pat[0]
|
|
add_q = None
|
|
reshape_q = pat[2]
|
|
mm_k = pat[3]
|
|
add_k = None
|
|
reshape_k = pat[5]
|
|
mm_v = pat[6]
|
|
add_v = None
|
|
reshape_v = pat[8]
|
|
|
|
head_dim = [
|
|
reshape_q.result(0).shape[-1],
|
|
reshape_k.result(0).shape[-1],
|
|
reshape_v.result(0).shape[-1],
|
|
]
|
|
fusion_w_name = f"fused_{mm_q.operand_source(1).name}_{mm_k.operand_source(1).name}_{mm_v.operand_source(1).name}"
|
|
fusion_map["qkv"].append(
|
|
{
|
|
fusion_w_name: [
|
|
ParamMeta(
|
|
mm_q.operand_source(1).name,
|
|
None,
|
|
None,
|
|
reshape_q.result(0).shape[-1],
|
|
),
|
|
ParamMeta(
|
|
mm_k.operand_source(1).name,
|
|
None,
|
|
None,
|
|
reshape_k.result(0).shape[-1],
|
|
),
|
|
ParamMeta(
|
|
mm_v.operand_source(1).name,
|
|
None,
|
|
None,
|
|
reshape_v.result(0).shape[-1],
|
|
),
|
|
]
|
|
}
|
|
)
|
|
fusion_w_dtype = mm_q.operand_source(1).dtype
|
|
fusion_w_shape = mm_q.operand_source(1).shape
|
|
fusion_w_shape[-1] += (
|
|
mm_k.operand_source(1).shape[-1] + mm_v.operand_source(1).shape[-1]
|
|
)
|
|
fusion_w_process_mesh = mm_q.operand_source(1).process_mesh
|
|
# insert fusion parameter
|
|
with paddle.static.program_guard(main_program, startup_program):
|
|
fused_w = paddle.pir.core.create_parameter(
|
|
dtype=fusion_w_dtype,
|
|
shape=fusion_w_shape,
|
|
name=fusion_w_name,
|
|
process_mesh=fusion_w_process_mesh,
|
|
placements=[
|
|
paddle.distributed.Replicate(),
|
|
paddle.distributed.Shard(1),
|
|
],
|
|
initializer=paddle.nn.initializer.Constant(value=0),
|
|
)
|
|
name2pir_param_map[fusion_w_name] = fused_w
|
|
if add_q is not None and add_k is not None and add_v is not None:
|
|
fusion_bias_name = f"fused_{add_q.operand_source(1).name}_{add_k.operand_source(1).name}_{add_v.operand_source(1).name}"
|
|
fusion_map["qkv"].append(
|
|
{
|
|
fusion_bias_name: [
|
|
ParamMeta(
|
|
add_q.operand_source(1).name,
|
|
None,
|
|
None,
|
|
reshape_q.result(0).shape[-1],
|
|
),
|
|
ParamMeta(
|
|
add_k.operand_source(1).name,
|
|
None,
|
|
None,
|
|
reshape_k.result(0).shape[-1],
|
|
),
|
|
ParamMeta(
|
|
add_v.operand_source(1).name,
|
|
None,
|
|
None,
|
|
reshape_v.result(0).shape[-1],
|
|
),
|
|
]
|
|
}
|
|
)
|
|
fusion_bias_dtype = add_q.operand_source(1).dtype
|
|
fusion_bias_shape = add_q.operand_source(1).shape
|
|
fusion_bias_shape[-1] += (
|
|
add_k.operand_source(1).shape[-1]
|
|
+ add_v.operand_source(1).shape[-1]
|
|
)
|
|
fusion_bias_process_mesh = add_q.operand_source(1).process_mesh
|
|
# insert fusion parameter
|
|
with paddle.static.program_guard(main_program, startup_program):
|
|
fused_bias = paddle.pir.core.create_parameter(
|
|
dtype=fusion_bias_dtype,
|
|
shape=fusion_bias_shape,
|
|
name=fusion_bias_name,
|
|
process_mesh=fusion_bias_process_mesh,
|
|
placements=[
|
|
paddle.distributed.Replicate(),
|
|
paddle.distributed.Shard(0),
|
|
],
|
|
initializer=paddle.nn.initializer.Constant(value=0),
|
|
)
|
|
name2pir_param_map[fusion_bias_name] = fused_bias
|
|
# insert dst pattern
|
|
paddle.pir.set_insertion_point_after(pat[-1])
|
|
fused_o = paddle.matmul(
|
|
mm_q.operand_source(0),
|
|
fused_w,
|
|
transpose_x=False,
|
|
transpose_y=False,
|
|
)
|
|
fused_o.get_defining_op().copy_attrs_from(mm_q)
|
|
if add_q is not None and add_k is not None and add_v is not None:
|
|
fused_o = paddle.add(fused_o, fused_bias)
|
|
fused_o.get_defining_op().copy_attrs_from(add_q)
|
|
out = paddle.reshape(
|
|
fused_o,
|
|
shape=[
|
|
0,
|
|
0,
|
|
reshape_k.result(0).shape[-2],
|
|
int(
|
|
(
|
|
reshape_q.result(0).shape[-2]
|
|
/ reshape_k.result(0).shape[-2]
|
|
+ 2
|
|
)
|
|
* reshape_q.result(0).shape[-1]
|
|
),
|
|
],
|
|
)
|
|
out.get_defining_op().copy_attrs_from(reshape_q)
|
|
reshape_op = out.get_defining_op()
|
|
if reshape_op.has_attr("struct_name"):
|
|
full_int_array_op = reshape_op.operand_source(1).get_defining_op()
|
|
full_int_array_op.set_str_attr(
|
|
"struct_name", reshape_op.attrs()["struct_name"]
|
|
)
|
|
out_q, out_k, out_v = paddle.split(
|
|
out,
|
|
num_or_sections=[
|
|
int(
|
|
(
|
|
reshape_q.result(0).shape[-2]
|
|
/ reshape_k.result(0).shape[-2]
|
|
)
|
|
* reshape_q.result(0).shape[-1]
|
|
),
|
|
reshape_k.result(0).shape[-1],
|
|
reshape_v.result(0).shape[-1],
|
|
],
|
|
axis=-1,
|
|
)
|
|
if reshape_op.has_attr("struct_name"):
|
|
builtin_split_op = out_q.get_defining_op()
|
|
split_op = builtin_split_op.operand_source(0).get_defining_op()
|
|
builtin_split_op.set_str_attr(
|
|
"struct_name", reshape_op.attrs()["struct_name"]
|
|
)
|
|
split_op.set_str_attr(
|
|
"struct_name", reshape_op.attrs()["struct_name"]
|
|
)
|
|
full_int_array_op = split_op.operand_source(1).get_defining_op()
|
|
full_int_array_op.set_str_attr(
|
|
"struct_name", reshape_op.attrs()["struct_name"]
|
|
)
|
|
full_op = split_op.operand_source(2).get_defining_op()
|
|
full_op.set_str_attr(
|
|
"struct_name", reshape_op.attrs()["struct_name"]
|
|
)
|
|
if reshape_q.result(0).shape[-2] != reshape_k.result(0).shape[-2]:
|
|
out_q = paddle.reshape(
|
|
out_q,
|
|
shape=[
|
|
0,
|
|
0,
|
|
reshape_q.result(0).shape[-2],
|
|
reshape_q.result(0).shape[-1],
|
|
],
|
|
)
|
|
if builtin_split_op.has_attr("struct_name"):
|
|
reshape_op = out_q.get_defining_op()
|
|
reshape_op.set_str_attr(
|
|
"struct_name", builtin_split_op.attrs()["struct_name"]
|
|
)
|
|
full_int_array_op = reshape_op.operand_source(
|
|
1
|
|
).get_defining_op()
|
|
full_int_array_op.set_str_attr(
|
|
"struct_name", builtin_split_op.attrs()["struct_name"]
|
|
)
|
|
|
|
reshape_q.result(0).replace_all_uses_with(out_q)
|
|
reshape_k.result(0).replace_all_uses_with(out_k)
|
|
reshape_v.result(0).replace_all_uses_with(out_v)
|
|
|
|
# 3. Delete src pattern from origin program.
|
|
del_ops = []
|
|
for pat in fused_pattern_map['ffn']:
|
|
for op in reversed(pat):
|
|
del_ops.append(op)
|
|
if op.name() == "pd_op.matmul" or op.name() == "pd_op.add":
|
|
del_ops.append(op.operand_source(1).get_defining_op())
|
|
del_ops.extend(
|
|
get_param_op(startup_program, op.operand_source(1).name)
|
|
)
|
|
for pat in fused_pattern_map['qkv']:
|
|
for op in reversed(pat):
|
|
del_ops.append(op)
|
|
if op.name() == "pd_op.matmul" or op.name() == "pd_op.add":
|
|
del_ops.append(op.operand_source(1).get_defining_op())
|
|
del_ops.extend(
|
|
get_param_op(startup_program, op.operand_source(1).name)
|
|
)
|
|
for op in del_ops:
|
|
op.erase()
|
|
|
|
# 4. Initialize fused parameters and delete original parameters.
|
|
concated_dy_param_index = []
|
|
# for key, pat_list in fused_name_map.items():
|
|
for key, pat_list in fusion_map.items():
|
|
for pat in pat_list:
|
|
for pir_param, dy_param_list in pat.items():
|
|
# Retrieve the params of ffn and qkv patterns from concrete_program for fusion.
|
|
concated_dy_param_list = []
|
|
for dy_param in dy_param_list:
|
|
param_index = dy_param_names.index(dy_param.name)
|
|
concated_dy_param_list.append(
|
|
concrete_program.parameters[0][param_index]
|
|
)
|
|
dy_param.local_shape = (
|
|
concrete_program.parameters[0][param_index]
|
|
._local_value()
|
|
.shape
|
|
)
|
|
if dy_param.local_head_dims is not None:
|
|
dy_param.local_num_head = (
|
|
dy_param.local_shape[-1] // dy_param.local_head_dims
|
|
)
|
|
concated_dy_param_index.append(param_index)
|
|
|
|
dy_param_init = True
|
|
for p in concated_dy_param_list:
|
|
if not p._local_value()._is_initialized():
|
|
dy_param_init = False
|
|
break
|
|
|
|
# Fuse params and init pir program fusion params.
|
|
with paddle.base.dygraph.guard():
|
|
dyparam_dtype = concated_dy_param_list[0].dtype
|
|
for param in concated_dy_param_list:
|
|
assert dyparam_dtype == param.dtype, (
|
|
"The dtypes of dy parameters to be fused are not the same."
|
|
)
|
|
|
|
dtensor = paddle.zeros(
|
|
shape=name2pir_param_map[pir_param].shape,
|
|
dtype=dyparam_dtype,
|
|
)
|
|
fused_dy_param = EagerParamBase.from_tensor(dtensor)
|
|
fused_dy_param = dist.shard_tensor(
|
|
fused_dy_param,
|
|
concated_dy_param_list[0].process_mesh,
|
|
concated_dy_param_list[0].placements,
|
|
)
|
|
fused_dy_param.name = pir_param
|
|
|
|
if dy_param_init:
|
|
if len(dy_param_list) == 3:
|
|
is_qkv = True
|
|
num_heads = dy_param_list[0].local_num_head
|
|
num_key_value_heads = dy_param_list[
|
|
1
|
|
].local_num_head
|
|
else:
|
|
is_qkv = False
|
|
num_heads = None
|
|
num_key_value_heads = None
|
|
concated_param = fuse_param_func(
|
|
[
|
|
obj._local_value()
|
|
for obj in concated_dy_param_list
|
|
],
|
|
is_qkv=is_qkv,
|
|
num_heads=num_heads,
|
|
num_key_value_heads=num_key_value_heads,
|
|
)
|
|
paddle.assign(
|
|
concated_param, fused_dy_param._local_value()
|
|
)
|
|
concated_param._clear()
|
|
|
|
# Pop and release original params from concrete_program
|
|
for param in concated_dy_param_list:
|
|
param.get_tensor()._clear()
|
|
|
|
concrete_program.parameters[0].append(fused_dy_param)
|
|
concrete_program.parameters[1].append(
|
|
name2pir_param_map[pir_param]
|
|
)
|
|
|
|
concated_dy_param_index.sort(reverse=True)
|
|
for index in concated_dy_param_index:
|
|
concrete_program.parameters[0].pop(index)
|
|
concrete_program.parameters[1].pop(index)
|
|
|
|
return fusion_map
|