chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,145 @@
|
||||
# Copyright (c) 2021 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.
|
||||
|
||||
from .allreduce_matmul_grad_overlapping import ( # noqa: F401
|
||||
AllreduceMatmulGradOverlappingPass,
|
||||
)
|
||||
from .auto_parallel_amp import ( # noqa: F401
|
||||
AMPLists,
|
||||
AMPPass,
|
||||
AMPState,
|
||||
)
|
||||
from .auto_parallel_c_embedding import ( # noqa: F401
|
||||
AutoParallelCEmbeddingPass,
|
||||
)
|
||||
from .auto_parallel_data_parallel_optimization import ( # noqa: F401
|
||||
DataParallelOptimizationPass,
|
||||
GradientsGroup,
|
||||
)
|
||||
from .auto_parallel_fp16 import ( # noqa: F401
|
||||
FP16Pass,
|
||||
FP16State,
|
||||
cast_startup_program,
|
||||
set_auto_cast_attr,
|
||||
set_op_dtype_to_fp16,
|
||||
)
|
||||
from .auto_parallel_fused_linear_promotion import ( # noqa: F401
|
||||
FusedLinearPromotionPass,
|
||||
)
|
||||
from .auto_parallel_grad_clip import ( # noqa: F401
|
||||
ClipGradByGlobalNormPass,
|
||||
ClipHelper,
|
||||
)
|
||||
from .auto_parallel_gradient_merge import ( # noqa: F401
|
||||
GradientMergePass,
|
||||
)
|
||||
from .auto_parallel_master_grad import ( # noqa: F401
|
||||
MasterGradPass,
|
||||
get_output_in_varlist,
|
||||
)
|
||||
from .auto_parallel_quantization import QuantizationPass # noqa: F401
|
||||
from .auto_parallel_recompute import ( # noqa: F401
|
||||
RecomputePass,
|
||||
RecomputeState,
|
||||
)
|
||||
from .auto_parallel_recompute_pir import ( # noqa: F401
|
||||
AutoParallelRecomputePIRPass,
|
||||
)
|
||||
from .auto_parallel_replace_with_parallel_cross_entropy import ( # noqa: F401
|
||||
AutoParallelReplaceWithParallelCrossEntropyPass,
|
||||
)
|
||||
from .auto_parallel_sequence_parallel_optimization import ( # noqa: F401
|
||||
SequenceParallelOptimizationPass,
|
||||
)
|
||||
from .auto_parallel_sharding import ( # noqa: F401
|
||||
ShardingInfo,
|
||||
ShardingPass,
|
||||
VarGroup,
|
||||
group_param,
|
||||
is_sharding_param_broadcast_op,
|
||||
partition_by_greedy_even,
|
||||
partition_by_use_order,
|
||||
partition_parameters,
|
||||
re_order_program,
|
||||
)
|
||||
from .auto_parallel_supplement_explicit_dependencies import ( # noqa: F401
|
||||
AutoParalSupplementDepPass,
|
||||
)
|
||||
from .auto_parallel_sync_shared_params import ( # noqa: F401
|
||||
AutoParallelSyncSharedParamsPass,
|
||||
)
|
||||
from .cpp_pass import ( # noqa: F401
|
||||
BuildCINNPass,
|
||||
FuseAdamWPass,
|
||||
FuseBatchNormActPass,
|
||||
FuseBatchNormAddActPass,
|
||||
FusedAttentionPass,
|
||||
FusedFeedforwardPass,
|
||||
FuseDotProductAttentionPass,
|
||||
FuseElementwiseAddActPass,
|
||||
FuseGemmEpiloguePass,
|
||||
FuseOptimizerPass,
|
||||
FuseReluDepthwiseConvPass,
|
||||
FuseResUnitPass,
|
||||
)
|
||||
|
||||
# InplaceAddtoOpPass,
|
||||
from .fuse_all_reduce import ( # noqa: F401
|
||||
FuseAllReducePass,
|
||||
filter_all_collective_op_indices,
|
||||
find_adjacent_match_sequences,
|
||||
find_all_fuse_all_reduce_groups,
|
||||
has_same_attrs,
|
||||
insert_coalesce_tensor_ops,
|
||||
insert_fuse_all_reduce_by_memory_size,
|
||||
insert_fuse_all_reduce_ops,
|
||||
split_fuse_all_reduce_groups_by_deps,
|
||||
)
|
||||
from .pass_base import PassContext, PassManager, new_pass
|
||||
from .pipeline_scheduler_pass import ( # noqa: F401
|
||||
Pipeline1F1BPass,
|
||||
PipelineEager1F1BPass,
|
||||
PipelineFThenBPass,
|
||||
PipelineVirtualPipelinePass,
|
||||
PipelineZeroBubblePipelinePass,
|
||||
apply_pass,
|
||||
)
|
||||
from .ps_server_pass import ( # noqa: F401
|
||||
AddGeoOptimizerPass,
|
||||
AddListenAndServPass,
|
||||
AddLrDecayTablePass,
|
||||
AddOptimizerPass,
|
||||
AddRpcGlobalFlagsPass,
|
||||
BuildPserverStartupProgramPass,
|
||||
DeleteUnusedInStartupPass,
|
||||
)
|
||||
from .ps_trainer_pass import ( # noqa: F401
|
||||
AppendSendOpsPass,
|
||||
DeleteExtraOptimizerPass,
|
||||
DeleteOptimizesPass,
|
||||
DistributedOpsPass,
|
||||
FakeInitOpsPass,
|
||||
PsGpuPass,
|
||||
PsTranspilePass,
|
||||
SetHeterPipelineOptPass,
|
||||
SplitFlOpsPass,
|
||||
SplitHeterWorkerOpsPass,
|
||||
SplitTrainerOpsPass,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
'new_pass',
|
||||
'PassManager',
|
||||
'PassContext',
|
||||
]
|
||||
@@ -0,0 +1,160 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import collections
|
||||
import logging
|
||||
|
||||
from ..auto_parallel.static.utils import (
|
||||
get_logger,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
from .pass_utils import AutoParallelStreamType, split_matmul_grad_to_matmul
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
# For allreduce pattern in the backward phase of column parallel linear:
|
||||
# dX, dY = matmul_grad(X, Y, dOut)
|
||||
# dX = all_reduce_sum(dX)
|
||||
# Split matmul_grad to 2 matmul:
|
||||
# dX = matmul(dOut, Y^T)
|
||||
# dX = all_reduce_sum(dX)
|
||||
# dY = matmul(X^T, dOut)
|
||||
#
|
||||
# Then the all_reduce sum can overlap with the compute of dY.
|
||||
@register_pass("allreduce_matmul_grad_overlapping")
|
||||
class AllreduceMatmulGradOverlappingPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.op_namescope = "/auto_parallel/allreduce_matmul_grad_overlapping"
|
||||
self.set_attr("dist_context", None)
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
self.dist_context = self.get_attr("dist_context")
|
||||
block = main_program.global_block()
|
||||
|
||||
matmul_grad_id_to_allreduce_id = (
|
||||
self._get_all_matmul_grad_and_allreduce_pairs(block)
|
||||
)
|
||||
logger.info(
|
||||
f"overlap matmul_grad and allreduce: {matmul_grad_id_to_allreduce_id}"
|
||||
)
|
||||
|
||||
self._split_matmul_grad_and_multi_streaming_allreduce(
|
||||
block, matmul_grad_id_to_allreduce_id
|
||||
)
|
||||
|
||||
def _get_all_matmul_grad_and_allreduce_pairs(self, block):
|
||||
ops = block.ops
|
||||
op_num = len(ops)
|
||||
matmul_grad_id_to_allreduce_id = collections.OrderedDict()
|
||||
for i, op_i in enumerate(ops):
|
||||
if (
|
||||
op_i.type == 'matmul_v2_grad'
|
||||
and op_i.attr("trans_x") is False
|
||||
and op_i.attr("trans_y") is False
|
||||
):
|
||||
x_grad = op_i.output("X@GRAD")
|
||||
for j in range(i + 1, op_num):
|
||||
op_j = ops[j]
|
||||
if (
|
||||
op_j.type == 'c_allreduce_sum'
|
||||
and op_j.input("X") == x_grad
|
||||
):
|
||||
matmul_grad_id_to_allreduce_id[i] = j
|
||||
return matmul_grad_id_to_allreduce_id
|
||||
|
||||
def _split_matmul_grad_and_multi_streaming_allreduce(
|
||||
self, block, matmul_grad_id_to_allreduce_id
|
||||
):
|
||||
ops = block.ops
|
||||
|
||||
for matmul_grad_id, allreduce_id in reversed(
|
||||
matmul_grad_id_to_allreduce_id.items()
|
||||
):
|
||||
matmul_grad_op = ops[matmul_grad_id]
|
||||
allreduce_op = ops[allreduce_id]
|
||||
|
||||
# NOTE(Sonder): When there are ops between matmul_grad and allreduce, we should check whether
|
||||
# these ops rely on the output of the intermediate ops. If so, we should not split the matmul_grad.
|
||||
# Otherwise, the output of the intermediate ops will get wrong results.
|
||||
skip_overlapping = False
|
||||
moved_ops_output = []
|
||||
matmul_grad_output = matmul_grad_op.output('Y@GRAD')[0]
|
||||
|
||||
for idx in range(matmul_grad_id + 1, allreduce_id):
|
||||
if matmul_grad_output in ops[idx].desc.input_arg_names():
|
||||
moved_ops_output.extend(ops[idx].desc.output_arg_names())
|
||||
else:
|
||||
for input_name in ops[idx].desc.input_arg_names():
|
||||
if input_name in moved_ops_output:
|
||||
skip_overlapping = True
|
||||
|
||||
if skip_overlapping:
|
||||
continue
|
||||
|
||||
# matmul_grad_op => matmul_v2 + reshape + reshape + matmul_v2 + reshape
|
||||
split_matmul_grad_to_matmul(
|
||||
block, matmul_grad_id, self.dist_context, self.op_namescope
|
||||
)
|
||||
|
||||
# NOTE(Ruibiao): Required OP scheduling order: matmul(dOut, Y^T) -> all_reduce_sum(dX) -> matmul(X^T, dOut).
|
||||
# all_reduce_sum(dX) and matmul(X^T, dOut) cannot be swapped. Otherwise, after buffer_shared_inplace_pass
|
||||
# adding share_buffer OP before all_reduce_sum, all_reduce_sum will synchronous with comp-stream, and then
|
||||
# the matmul op before it cannot be overlapped.
|
||||
allreduce_op_dist_attr = (
|
||||
self.dist_context.get_op_dist_attr_for_program(allreduce_op)
|
||||
)
|
||||
allreduce_op_dist_attr.execution_stream = (
|
||||
AutoParallelStreamType.MP_STREAM.value
|
||||
)
|
||||
|
||||
allreduce_op_inputs = allreduce_op.desc.input_names()
|
||||
allreduce_op_outputs = allreduce_op.desc.output_names()
|
||||
|
||||
allreduce_op_inputs = {
|
||||
name: allreduce_op.input(name) for name in allreduce_op_inputs
|
||||
}
|
||||
allreduce_op_outputs = {
|
||||
name: allreduce_op.output(name) for name in allreduce_op_outputs
|
||||
}
|
||||
|
||||
# matmul_v2 + reshape + reshape + matmul_v2 + reshape + ... + original all_reduce_sum
|
||||
# =>
|
||||
# matmul_v2 + new all_reduce_sum + reshape + reshape + matmul_v2 + reshape + ... + original all_reduce_sum
|
||||
#
|
||||
# NOTE(liym27): new all_reduce_sum must be inserted to "the next of the first matmul_v2", otherwise another
|
||||
# pass fused_linear_param_grad_add will not work.
|
||||
allreduce_op = block._insert_op_without_sync(
|
||||
index=matmul_grad_id + 1,
|
||||
type=allreduce_op.type,
|
||||
inputs=allreduce_op_inputs,
|
||||
outputs=allreduce_op_outputs,
|
||||
attrs=allreduce_op.all_attrs(),
|
||||
)
|
||||
self.dist_context.set_op_dist_attr_for_program(
|
||||
allreduce_op, allreduce_op_dist_attr
|
||||
)
|
||||
# Remove the original allreduce op
|
||||
block._remove_op(allreduce_id + 5, sync=False)
|
||||
|
||||
block._sync_with_cpp()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,345 @@
|
||||
# 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 re
|
||||
import warnings
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.base.core import TensorDistAttr
|
||||
from paddle.distributed import fleet
|
||||
from paddle.distributed.auto_parallel.static.dist_attribute import (
|
||||
DistTensorSpec,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
from paddle.framework import core
|
||||
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
|
||||
@register_pass("auto_parallel_c_embedding_pass")
|
||||
class AutoParallelCEmbeddingPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
mp_size = hcg.get_model_parallel_world_size()
|
||||
if mp_size > 1:
|
||||
return True
|
||||
warnings.warn("c_embedding pass is only applicable to tnesor parallel.")
|
||||
return False
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
concrete_program = self.get_attr("concrete_program")
|
||||
ops = main_program.global_block().ops
|
||||
for i, op in enumerate(ops):
|
||||
if op.name() == 'pd_op.embedding':
|
||||
# update weight dims mapping
|
||||
mp_axis = self._update_weight(op, concrete_program)
|
||||
# update startup_program
|
||||
self._update_startup_program(startup_program, mp_axis)
|
||||
# replace embedding with c_embedding
|
||||
c_emb_op = self._replace_embedding_with_c_embedding(op)
|
||||
# insert allreduce reshard
|
||||
comm_op = self._insert_allreduce_reshard(c_emb_op)
|
||||
# update dims_mapping before c_embedding
|
||||
self._update_before_dims_mapping(c_emb_op)
|
||||
# update dims_mapping after c_embedding
|
||||
self._update_after_dims_mapping(comm_op)
|
||||
|
||||
def _update_weight(self, op, concrete_program):
|
||||
# update weight dims_mapping concrete_program
|
||||
placements = op.operand(1).source().placements
|
||||
dim_map, partial_status = dist.auto_parallel.placement_type.to_dim_map(
|
||||
placements, op.operand(1).source().ndim
|
||||
)
|
||||
# mp_axis is used to specify the axis for row parallel
|
||||
mp_axis = -1
|
||||
dim_map = [-1, -1]
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
mp_size = hcg.get_model_parallel_world_size()
|
||||
if mp_size > 1:
|
||||
strategy = fleet.DistributedStrategy()
|
||||
# get mp_axis from DistributedStrategy
|
||||
mp_axis = strategy.hybrid_configs['mp_degree']
|
||||
dim_map = [mp_axis, -1]
|
||||
dist_attr_w = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
op.operand(1).source().process_mesh,
|
||||
dim_map,
|
||||
partial_status,
|
||||
)
|
||||
dist_type_input0 = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
op.operand(1).source().type(), dist_attr_w
|
||||
)
|
||||
op.operand(1).source().set_type(dist_type_input0)
|
||||
# update c_embedding weight dynamic parameters
|
||||
dy_params = concrete_program.parameters[0]
|
||||
pattern = re.compile(r'embedding_.*\.w_0\.dist')
|
||||
for index, param in enumerate(dy_params):
|
||||
if pattern.match(param.name):
|
||||
var_dist_attr = TensorDistAttr()
|
||||
var_dist_attr.process_mesh = dist_attr_w.process_mesh
|
||||
var_dist_attr.dims_mapping = dist_attr_w.dims_mapping
|
||||
tmp = paddle.base.core.reshard(param, var_dist_attr)
|
||||
param.get_tensor()._share_data_with(tmp.get_tensor())
|
||||
return mp_axis
|
||||
|
||||
def _replace_embedding_with_c_embedding(self, op):
|
||||
paddle.pir.set_insertion_point(op)
|
||||
num_embeddings = op.operand(1).source().type().shape[0]
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
# compute the start_index using the MP's world size and rank
|
||||
mp_size = hcg.get_model_parallel_world_size()
|
||||
mp_rank = hcg.get_model_parallel_rank()
|
||||
per_part_size = num_embeddings // mp_size
|
||||
vocab_start_index = mp_rank * per_part_size
|
||||
t_op = paddle._C_ops.c_embedding(
|
||||
op.operand(1).source(),
|
||||
op.operand(0).source(),
|
||||
vocab_start_index,
|
||||
num_embeddings,
|
||||
)
|
||||
t_op.get_defining_op().op_role = int(OpRole.Forward)
|
||||
new_op = t_op.get_defining_op()
|
||||
op.result(0).replace_all_uses_with(t_op)
|
||||
op.erase()
|
||||
return new_op
|
||||
|
||||
def _insert_allreduce_reshard(self, c_emb_op):
|
||||
result = c_emb_op.result(0)
|
||||
paddle.pir.set_insertion_point_after(c_emb_op)
|
||||
placements = result.dist_attr().placements
|
||||
dim_map, partial_status = dist.auto_parallel.placement_type.to_dim_map(
|
||||
placements, result.ndim
|
||||
)
|
||||
partial_status = {}
|
||||
dist_attr_new = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
result.process_mesh,
|
||||
dim_map,
|
||||
partial_status,
|
||||
)
|
||||
# insert allreduce by inserting reshard with an empty partial.
|
||||
comm_op_t = paddle._C_ops.reshard_v2(result, dist_attr_new)
|
||||
comm_op_t.get_defining_op().op_role = int(OpRole.Forward)
|
||||
result.replace_all_uses_with(comm_op_t)
|
||||
comm_op = comm_op_t.get_defining_op()
|
||||
comm_op.operand(0).set_source(result)
|
||||
return comm_op
|
||||
|
||||
def _update_before_dims_mapping(self, new_op):
|
||||
placements = new_op.operand(0).source().placements
|
||||
stack = [new_op.operand(0).source().get_defining_op()]
|
||||
# adjust all ops before c_embedding until parameters input
|
||||
while stack:
|
||||
op = stack.pop()
|
||||
operands, results = [], []
|
||||
if op.num_results() > 0:
|
||||
for result, result_dist in zip(
|
||||
op.results(), op.dist_attr.results()
|
||||
):
|
||||
placements_dist = (
|
||||
result_dist.as_tensor_dist_attr().placements
|
||||
)
|
||||
if placements != placements_dist:
|
||||
dim_map, partial_status = (
|
||||
dist.auto_parallel.placement_type.to_dim_map(
|
||||
placements, result.ndim
|
||||
)
|
||||
)
|
||||
dist_attr_new = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
result.process_mesh,
|
||||
dim_map,
|
||||
partial_status,
|
||||
)
|
||||
dist_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
result.type(), dist_attr_new
|
||||
)
|
||||
result.set_type(dist_type)
|
||||
results.append(dist_attr_new)
|
||||
sub_name = op.name().split('.')[1]
|
||||
if op.num_operands() > 0:
|
||||
assert sub_name != "cast", (
|
||||
"Need to add support for {sub_name}."
|
||||
)
|
||||
operands.append(dist_attr_new)
|
||||
next_op = op.operand(0).source().get_defining_op()
|
||||
stack.append(next_op)
|
||||
process_mesh = (
|
||||
op.results()[0].process_mesh
|
||||
if op.num_results() > 0
|
||||
else op.operand(0).source().process_mesh
|
||||
)
|
||||
op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
process_mesh,
|
||||
operands,
|
||||
results,
|
||||
)
|
||||
)
|
||||
|
||||
def _update_after_dims_mapping(self, new_op):
|
||||
placements = new_op.result(0).placements
|
||||
pre_id = new_op.id()
|
||||
stack = list(new_op.result(0).all_used_ops())
|
||||
# adjust all ops after c_embedding until the placements are consistent
|
||||
while stack:
|
||||
op = stack.pop()
|
||||
operands, results = [], []
|
||||
if op.num_operands() > 0:
|
||||
for operand, operand_dist in zip(
|
||||
op.operands_source(), op.dist_attr.operands()
|
||||
):
|
||||
if operand.get_defining_op().id() != pre_id:
|
||||
continue
|
||||
placements_dist = (
|
||||
operand_dist.as_tensor_dist_attr().placements
|
||||
)
|
||||
if placements != placements_dist:
|
||||
dim_map, partial_status = (
|
||||
dist.auto_parallel.placement_type.to_dim_map(
|
||||
placements, operand.ndim
|
||||
)
|
||||
)
|
||||
dist_attr_new = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
operand.process_mesh,
|
||||
dim_map,
|
||||
partial_status,
|
||||
)
|
||||
dist_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
operand.type(), dist_attr_new
|
||||
)
|
||||
operand.set_type(dist_type)
|
||||
operands.append(dist_attr_new)
|
||||
sub_name = op.name().split('.')[1]
|
||||
if sub_name == 'reshard':
|
||||
# only change reshard‘s inputs
|
||||
placements_out0 = op.results()[0].placements
|
||||
dim_map_out0, partial_status_out0 = (
|
||||
dist.auto_parallel.placement_type.to_dim_map(
|
||||
placements_out0,
|
||||
op.results()[0].ndim,
|
||||
)
|
||||
)
|
||||
dist_attr_out0 = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
op.results()[0].process_mesh,
|
||||
dim_map_out0,
|
||||
partial_status_out0,
|
||||
)
|
||||
results.append(dist_attr_out0)
|
||||
elif core.contains_spmd_rule(sub_name):
|
||||
# redo the infer spmd_rule
|
||||
rule = core.get_phi_spmd_rule(sub_name)
|
||||
tensor_dist_attr = TensorDistAttr()
|
||||
tensor_dist_attr.dims_mapping = dim_map
|
||||
partial_dims = []
|
||||
for i, p in enumerate(placements):
|
||||
if isinstance(p, dist.Partial):
|
||||
partial_dims.append(i)
|
||||
if len(partial_dims) > 0:
|
||||
tensor_dist_attr._set_partial_dims(partial_dims)
|
||||
tensor_dist_attr.process_mesh = operand.process_mesh
|
||||
inputs = DistTensorSpec(
|
||||
operand.shape, tensor_dist_attr
|
||||
)
|
||||
attr_names = op.get_attr_names()
|
||||
input_specs = []
|
||||
input_specs.append(inputs)
|
||||
for attr_name in attr_names:
|
||||
input_specs.append(op.attrs()[attr_name])
|
||||
inferred_dist_attrs = rule.infer_forward(
|
||||
*input_specs
|
||||
)
|
||||
dims_mapping_new_out = inferred_dist_attrs[1][
|
||||
0
|
||||
].dims_mapping
|
||||
partial_status = {}
|
||||
if inferred_dist_attrs[1][0]._is_partial():
|
||||
partial_dims = inferred_dist_attrs[1][
|
||||
0
|
||||
]._partial_dims()
|
||||
for i in partial_dims:
|
||||
partial_status[i] = (
|
||||
paddle.base.core.ReduceType.kRedSum
|
||||
)
|
||||
dist_attr_new_out = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
operand.process_mesh,
|
||||
dims_mapping_new_out,
|
||||
partial_status,
|
||||
)
|
||||
dist_type = (
|
||||
paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
op.result(0).type(), dist_attr_new_out
|
||||
)
|
||||
)
|
||||
op.result(0).set_type(dist_type)
|
||||
results.append(dist_attr_new_out)
|
||||
next_op = op.results()[0].all_used_ops()[0]
|
||||
stack.append(next_op)
|
||||
pre_id = op.id()
|
||||
placements = dist_attr_new_out.placements
|
||||
else:
|
||||
results.append(dist_attr_new)
|
||||
next_op = op.results()[0].all_used_ops()[0]
|
||||
stack.append(next_op)
|
||||
pre_id = op.id()
|
||||
|
||||
process_mesh = (
|
||||
op.results()[0].process_mesh
|
||||
if op.num_results() > 0
|
||||
else op.operand(0).source().process_mesh
|
||||
)
|
||||
op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
process_mesh,
|
||||
operands,
|
||||
results,
|
||||
)
|
||||
)
|
||||
|
||||
def _update_startup_program(self, startup_program, mp_axis):
|
||||
# modify the startup_program because the optimizer needs to use
|
||||
startup_block = startup_program.global_block()
|
||||
for op in startup_block.ops:
|
||||
if op.name() == 'pd_op.full':
|
||||
next_op = op.result(0).all_used_ops()[0]
|
||||
parameter_name = next_op.str_attr("parameter_name")
|
||||
pattern = re.compile(r'embedding_.*\.w_0\.dist')
|
||||
if pattern.match(parameter_name):
|
||||
placements = op.results()[0].placements
|
||||
dim_map, partial_status = (
|
||||
dist.auto_parallel.placement_type.to_dim_map(
|
||||
placements, len(placements)
|
||||
)
|
||||
)
|
||||
dim_map = [mp_axis, -1]
|
||||
dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
op.results()[0].process_mesh,
|
||||
dim_map,
|
||||
partial_status,
|
||||
)
|
||||
)
|
||||
dist_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
op.results()[0].type(), dist_attr
|
||||
)
|
||||
op.results()[0].set_type(dist_type)
|
||||
op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
op.results()[0].process_mesh, [], [dist_attr]
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,766 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.process_mesh import ProcessMesh
|
||||
from paddle.distributed.auto_parallel.static.dist_attribute import (
|
||||
OperatorDistAttr,
|
||||
TensorDistAttr,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.operators.common import (
|
||||
is_data_parallel_reduce_op,
|
||||
is_data_parallel_scale_op,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
find_higher_order_backward_op,
|
||||
get_var_numel,
|
||||
insert_dependencies_for_vars,
|
||||
is_forward_op,
|
||||
is_loss_grad_op,
|
||||
is_optimize_op,
|
||||
ring_id_to_process_group,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
from paddle.static import default_main_program
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from .pass_base import PassBase, PassType, register_pass
|
||||
|
||||
# add new optimizers supporting rescale_grad here
|
||||
__rescale_grad_supported_opts__ = [
|
||||
'lars_momentum',
|
||||
'sparse_momentum',
|
||||
'dgc_momentum',
|
||||
'momentum',
|
||||
'merge_momentum',
|
||||
]
|
||||
|
||||
# a heuristic number
|
||||
__max_stream_num_allow__ = 16
|
||||
|
||||
|
||||
@register_pass("auto_parallel_data_parallel_optimization")
|
||||
class DataParallelOptimizationPass(PassBase):
|
||||
"""
|
||||
Apply Optimizations that specialized for data parallelism in Auto Parallel.
|
||||
1. prune grad scaling
|
||||
2. overlap comm and calc
|
||||
3. fuse allreduce
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# NOTE not use dependence on loss and param_grads
|
||||
self.set_attr("dist_context", None)
|
||||
self.set_attr("global_rank", -1)
|
||||
self.set_attr("use_sharding", False)
|
||||
# {grad1: group1, grad2: group1, grad3: group2}
|
||||
# record the order for fuse grad data memory
|
||||
self._grad_name_to_group_map = OrderedDict()
|
||||
# {group1:[grad1, grad2] , group2:[grad3]}
|
||||
self._group_to_grad_name_map = OrderedDict()
|
||||
self._support_rescale_grad = False
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
if (not isinstance(self.get_attr("global_rank"), int)) or self.get_attr(
|
||||
"global_rank"
|
||||
) < 0:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _type(self):
|
||||
return PassType.COMM_OPT
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
self.dist_context = self.get_attr("dist_context")
|
||||
self.global_rank = int(self.get_attr("global_rank"))
|
||||
self.use_sharding = self.get_attr("use_sharding")
|
||||
self.coalesce_prefix = 'coalesce_grad'
|
||||
self.gradient_sync_stream = "gradient_sync_stream"
|
||||
|
||||
with paddle.static.program_guard(main_program, startup_program):
|
||||
self._analyze_program()
|
||||
|
||||
# TODO refactor here to first fuse then overlap
|
||||
if self.is_data_parallel_applied():
|
||||
self._prune_grad_scaling()
|
||||
self._calc_comm_overlap()
|
||||
grad_group = self._fuse_allreduce()
|
||||
self._add_dependencies(grad_group)
|
||||
self.summary(grad_group)
|
||||
|
||||
def _prune_grad_scaling(self):
|
||||
if not self._could_be_prune():
|
||||
return
|
||||
|
||||
if self._all_dp_groups_same_degree():
|
||||
self._scale_backward_initial_grad()
|
||||
else:
|
||||
self._update_opt_rescale_grad()
|
||||
|
||||
self._remove_grad_scaling()
|
||||
|
||||
def _calc_comm_overlap(self):
|
||||
if not self._could_be_overlap():
|
||||
return
|
||||
self._comms_overlap_calc()
|
||||
self._calc_wait_comms()
|
||||
|
||||
def _fuse_allreduce(self):
|
||||
if not self._could_be_fuse():
|
||||
return []
|
||||
|
||||
grad_group = self._group_grads()
|
||||
self._update_program(grad_group)
|
||||
|
||||
return grad_group
|
||||
|
||||
def _analyze_program(self):
|
||||
"""
|
||||
build two maps
|
||||
{param_grad_name: data_parallel_group}
|
||||
{pdata_parallel_group: aram_grad_name}
|
||||
"""
|
||||
|
||||
block = default_main_program().global_block()
|
||||
ops = block.ops
|
||||
scaled_grads = []
|
||||
|
||||
for op in ops:
|
||||
if is_data_parallel_reduce_op(op):
|
||||
grad_name = op.output_arg_names[0]
|
||||
if grad_name in self._grad_name_to_group_map:
|
||||
continue
|
||||
assert op.has_attr("ring_id"), (
|
||||
f"Unexpected: comm op [{op}] has NOT ring id."
|
||||
)
|
||||
group = ring_id_to_process_group(op.attr("ring_id"))
|
||||
|
||||
assert group is not None, (
|
||||
f"Unexpected: data parallel group of [{grad_name}] from op [{op}] is None"
|
||||
)
|
||||
|
||||
self._grad_name_to_group_map[grad_name] = group
|
||||
|
||||
if group not in self._group_to_grad_name_map:
|
||||
self._group_to_grad_name_map[group] = [grad_name]
|
||||
else:
|
||||
self._group_to_grad_name_map[group].append(grad_name)
|
||||
|
||||
elif is_data_parallel_scale_op(op):
|
||||
grad_name = op.output_arg_names[0]
|
||||
scaled_grads.append(grad_name)
|
||||
|
||||
# TODO support multiple optimizers in on network in future.
|
||||
# here we assume that the optimizer is unique in network.
|
||||
elif (
|
||||
is_optimize_op(op)
|
||||
and op.type in __rescale_grad_supported_opts__
|
||||
):
|
||||
self._support_rescale_grad = True
|
||||
|
||||
not_synchronized_grads = []
|
||||
for grad_name in scaled_grads:
|
||||
if grad_name not in self._grad_name_to_group_map:
|
||||
not_synchronized_grads.append(grad_name)
|
||||
assert len(not_synchronized_grads) == 0, (
|
||||
f"Unexpected: gradients [{not_synchronized_grads}] is scaled BUT NOT synchronized."
|
||||
)
|
||||
|
||||
def is_data_parallel_applied(self):
|
||||
return len(self._group_to_grad_name_map) > 0
|
||||
|
||||
def _could_be_prune(self):
|
||||
return self.dist_context.gradient_scale and (
|
||||
self._support_rescale_grad or self._all_dp_groups_same_degree()
|
||||
)
|
||||
|
||||
def _all_dp_groups_same_degree(self):
|
||||
return (
|
||||
len(
|
||||
{
|
||||
len(group.ranks)
|
||||
for group in self._group_to_grad_name_map.keys()
|
||||
}
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
def _scale_backward_initial_grad(self):
|
||||
block = default_main_program().global_block()
|
||||
dp_degree = len(next(iter(self._group_to_grad_name_map.keys())).ranks)
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_loss_grad_op(op):
|
||||
assert op.type == 'fill_constant', (
|
||||
"loss_grad_op must be fill_constant op, "
|
||||
f"but this op is {op.type}"
|
||||
)
|
||||
assert op.has_attr('value')
|
||||
loss_scale = float(op.attr('value'))
|
||||
loss_scale = loss_scale / dp_degree
|
||||
op._set_attr('value', loss_scale)
|
||||
break
|
||||
|
||||
def _remove_grad_scaling(self):
|
||||
block = default_main_program().global_block()
|
||||
|
||||
for op_idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_data_parallel_scale_op(op):
|
||||
block._remove_op(op_idx, False)
|
||||
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _update_opt_rescale_grad(self):
|
||||
block = default_main_program().global_block()
|
||||
scaled_grads = set()
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if (
|
||||
is_optimize_op(op)
|
||||
and op.type in __rescale_grad_supported_opts__
|
||||
):
|
||||
assert op.has_attr('rescale_grad'), (
|
||||
f"Unexpected: op [{op}] is supported to have [rescale_grad] attribute."
|
||||
)
|
||||
assert len(op.input("Grad")) == 1, (
|
||||
f"Unexpected: op [{op}] is supported to have only one input grad var."
|
||||
)
|
||||
|
||||
grad_name = op.input("Grad")[0]
|
||||
dp_degree = len(
|
||||
list(self._grad_name_to_group_map[grad_name].ranks)
|
||||
)
|
||||
scaled_grads.add(grad_name)
|
||||
|
||||
rescale_grad = float(op.attr('rescale_grad')) / dp_degree
|
||||
op._set_attr('rescale_grad', rescale_grad)
|
||||
|
||||
assert scaled_grads == set(self._grad_name_to_group_map.keys()), (
|
||||
f"Unexpected: gradients [{set(self._grad_name_to_group_map.keys()) - scaled_grads}] are unscaled."
|
||||
)
|
||||
|
||||
def _could_be_overlap(self):
|
||||
# NOTE current different nccl comm will use different cuda stream
|
||||
# so if there too many dp group there will be too many stream need to be
|
||||
# created and sync.
|
||||
# revise here when framework support custom stream in static graph mode.
|
||||
num_dp_comm_stream = len(set(self._group_to_grad_name_map.keys()))
|
||||
if num_dp_comm_stream > __max_stream_num_allow__:
|
||||
return False
|
||||
if self.use_sharding:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _comms_overlap_calc(self):
|
||||
# TODO support InterpreterCore executor for overlap.
|
||||
# InterpreterCore has a different logic for overlapping
|
||||
# which is different from use_calc_stream
|
||||
block = default_main_program().global_block()
|
||||
|
||||
# comm wait calc to finish
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_data_parallel_reduce_op(op):
|
||||
assert op.has_attr('ring_id')
|
||||
|
||||
op._set_attr('use_calc_stream', False)
|
||||
ring_id = op.attr("ring_id")
|
||||
block._insert_op_without_sync(
|
||||
idx,
|
||||
type='c_wait_compute',
|
||||
inputs={'X': []},
|
||||
outputs={'Out': []},
|
||||
attrs={'op_role': OpRole.Backward, 'ring_id': ring_id},
|
||||
)
|
||||
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _calc_wait_comms(self):
|
||||
return
|
||||
|
||||
block = default_main_program().global_block()
|
||||
|
||||
# NOTE the naive overlap implement in static hybrid parallel only sync comm stream
|
||||
# at the end of Backward phase, based on a strong constraint that
|
||||
# all communicating gradient would NOT be used after communication in Backward phase.
|
||||
# BUT this constraint will fail for scenario like Weight-Sharing and Higher-Order Differentiation,
|
||||
# where gradient will be involved in other calculation between data-parallel allreduce kernel submitted
|
||||
# into comm streams and the synchronization of comm stream at the end of Backward phase.
|
||||
# synchronization of comm stream should add according to the usage of communicating gradients
|
||||
# to support Overlapping for Weight-Sharing and Higher-Order Differentiation.
|
||||
|
||||
ring_id_to_un_sync_grad_map = {}
|
||||
op_idx_to_sync_ring_id_map = {}
|
||||
for group in self._group_to_grad_name_map.keys():
|
||||
ring_id_to_un_sync_grad_map[group.id] = []
|
||||
|
||||
# analyze the where need to sync
|
||||
for i, op in enumerate(block.ops):
|
||||
if is_data_parallel_reduce_op(op):
|
||||
ring_id = op.attr("ring_id")
|
||||
grad_name = op.output_arg_names[0]
|
||||
ring_id_to_un_sync_grad_map[ring_id].append(grad_name)
|
||||
elif is_data_parallel_scale_op(op):
|
||||
continue
|
||||
# other ops that might use communicating grad
|
||||
else:
|
||||
for input_var_name in op.input_arg_names:
|
||||
for (
|
||||
ring_id,
|
||||
unsync_grad_names,
|
||||
) in ring_id_to_un_sync_grad_map.items():
|
||||
if input_var_name in unsync_grad_names:
|
||||
# need to sync before op_i
|
||||
if i in op_idx_to_sync_ring_id_map:
|
||||
op_idx_to_sync_ring_id_map[i].append(ring_id)
|
||||
else:
|
||||
op_idx_to_sync_ring_id_map[i] = [ring_id]
|
||||
# all grads in this comm stream are synced
|
||||
ring_id_to_un_sync_grad_map[ring_id] = []
|
||||
|
||||
# insert synchronization
|
||||
indices = list(op_idx_to_sync_ring_id_map.keys())
|
||||
# TODO the synchronization could be optimized
|
||||
# we should record the event of a gradient is communicating and
|
||||
# only wait for that event to be completed.
|
||||
# BUT paddle static currently not support op api for event record only, so
|
||||
# here we try to wait for all kernel in that comm stream to be finish which is not that optimized.
|
||||
for i in sorted(indices, reverse=True):
|
||||
for ring_id in op_idx_to_sync_ring_id_map[i]:
|
||||
block._insert_op_without_sync(
|
||||
i,
|
||||
type='c_wait_comm',
|
||||
inputs={'X': []},
|
||||
outputs={'Out': []},
|
||||
attrs={'op_role': OpRole.Backward, 'ring_id': ring_id},
|
||||
)
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _could_be_fuse(self):
|
||||
# TODO support gradient fuse higher order gradient.
|
||||
# should analyse the dependencies of gradient in backward.
|
||||
if find_higher_order_backward_op(default_main_program()):
|
||||
return False
|
||||
if self.use_sharding:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _group_grads(self):
|
||||
"""
|
||||
conditions for gradients to be grouped:
|
||||
1. group size < max_fuse_numel
|
||||
2. same dp group
|
||||
3. same dtype
|
||||
4. dependency: grad would NOT be used by other ops within group segment
|
||||
|
||||
gradients inside same group would be fuse into one coalesce tensor
|
||||
"""
|
||||
|
||||
block = default_main_program().global_block()
|
||||
ops = block.ops
|
||||
|
||||
# group individual grad vars
|
||||
# TODO consider fuse gradient for sharding reduce
|
||||
# TODO let user to set fuse_grad_size
|
||||
# emb = 50000 * h, ffn = 8 * h * h, mha = 4 * h * h
|
||||
h = 2048
|
||||
ffn_numel = 2 * (4 * h) * h
|
||||
mha_numel = 3 * h * h + h * h
|
||||
max_fuse_numel = ffn_numel + mha_numel
|
||||
grad_groups = []
|
||||
cur_group = GradientsGroup(ops, max_fuse_numel)
|
||||
grouped_grad_names = set()
|
||||
|
||||
def collect_group(cur_group, grad_var, ring_id, i):
|
||||
if len(cur_group.gradients) == 0:
|
||||
cur_group = None
|
||||
else:
|
||||
cur_group.finalize()
|
||||
grad_groups.append(cur_group)
|
||||
|
||||
new_group = GradientsGroup(ops, max_fuse_numel)
|
||||
if grad_var:
|
||||
new_group.add(grad_var, ring_id, i)
|
||||
grouped_grad_names.add(grad_var.name)
|
||||
return new_group
|
||||
|
||||
def op_depend_on_group(op, group):
|
||||
vars_ = set(op.input_arg_names + op.output_arg_names)
|
||||
grad_names = {grad.name for grad in group.gradients}
|
||||
return len(vars_.intersection(grad_names)) > 0
|
||||
|
||||
for i, op in enumerate(ops):
|
||||
if is_data_parallel_reduce_op(op):
|
||||
ring_id = op.attr("ring_id")
|
||||
grad_name = op.output_arg_names[0]
|
||||
grad_var = block.var(grad_name)
|
||||
grad_numel = get_var_numel(grad_var)
|
||||
|
||||
if cur_group.acceptable(grad_var, ring_id):
|
||||
assert grad_name not in grouped_grad_names
|
||||
grouped_grad_names.add(grad_name)
|
||||
cur_group.add(grad_var, ring_id, i)
|
||||
else:
|
||||
cur_group = collect_group(cur_group, grad_var, ring_id, i)
|
||||
else:
|
||||
if op_depend_on_group(op, cur_group):
|
||||
cur_group = collect_group(cur_group, None, None, None)
|
||||
|
||||
# collect last group
|
||||
collect_group(cur_group, None, None, None)
|
||||
|
||||
return grad_groups
|
||||
|
||||
def _update_program(self, grad_groups):
|
||||
block = default_main_program().global_block()
|
||||
|
||||
remove_op_types = [
|
||||
'scale',
|
||||
'all_reduce',
|
||||
'c_wait_compute',
|
||||
]
|
||||
|
||||
for i, group in enumerate(grad_groups[::-1]):
|
||||
# skip unfused big tensor
|
||||
if len(group.gradients) <= 1:
|
||||
group.coalesce_var = group.gradients[0]
|
||||
continue
|
||||
|
||||
ref_process_mesh = set()
|
||||
concated_shapes = []
|
||||
concated_ranks = []
|
||||
for grad_ in group.gradients:
|
||||
grad_dist_attr = (
|
||||
self.dist_context.get_tensor_dist_attr_for_program(grad_)
|
||||
)
|
||||
ref_process_mesh.update(
|
||||
set(grad_dist_attr.process_mesh.process_ids)
|
||||
)
|
||||
|
||||
shape = grad_.shape
|
||||
concated_shapes.extend(shape)
|
||||
concated_ranks.append(len(shape))
|
||||
|
||||
# create coalesce tensor
|
||||
group.coalesce_var = block.create_var(
|
||||
name=unique_name.generate(self.coalesce_prefix + f'_{i}'),
|
||||
dtype=group.dtype,
|
||||
persistable=False,
|
||||
stop_gradient=True,
|
||||
)
|
||||
|
||||
tensor_dist_attr = TensorDistAttr()
|
||||
tensor_dist_attr.process_mesh = ProcessMesh(list(ref_process_mesh))
|
||||
tensor_dist_attr.dims_mapping = []
|
||||
self.dist_context.set_tensor_dist_attr_for_program(
|
||||
group.coalesce_var, tensor_dist_attr
|
||||
)
|
||||
|
||||
# update allreduce & scale op
|
||||
if group.scale_op_idx != -1:
|
||||
scale_op = block.ops[group.scale_op_idx]
|
||||
assert scale_op.type == 'scale', (
|
||||
f"should found scale op but found {scale_op}"
|
||||
)
|
||||
scale_op._rename_input(
|
||||
scale_op.input_arg_names[0], group.coalesce_var.name
|
||||
)
|
||||
scale_op._rename_output(
|
||||
scale_op.output_arg_names[0], group.coalesce_var.name
|
||||
)
|
||||
|
||||
allreduce_op = block.ops[group.allreduce_op_idx]
|
||||
assert (
|
||||
allreduce_op.type == 'all_reduce'
|
||||
and allreduce_op.attr('reduce_type')
|
||||
== paddle.distributed.ReduceOp.SUM
|
||||
), f"should found all_reduce sum op but found {allreduce_op}"
|
||||
allreduce_op_dist_attr = (
|
||||
self.dist_context.get_op_dist_attr_for_program(allreduce_op)
|
||||
)
|
||||
old_in_name = allreduce_op.input_arg_names[0]
|
||||
new_in_name = group.coalesce_var.name
|
||||
allreduce_op._rename_input(old_in_name, new_in_name)
|
||||
input_dist_attr = allreduce_op_dist_attr.get_input_dist_attr(
|
||||
old_in_name
|
||||
)
|
||||
allreduce_op_dist_attr.set_input_dist_attr(
|
||||
new_in_name, input_dist_attr
|
||||
)
|
||||
|
||||
old_out_name = allreduce_op.output_arg_names[0]
|
||||
new_out_name = group.coalesce_var.name
|
||||
allreduce_op._rename_output(old_out_name, new_out_name)
|
||||
out_dist_attr = allreduce_op_dist_attr.get_output_dist_attr(
|
||||
old_out_name
|
||||
)
|
||||
allreduce_op_dist_attr.set_output_dist_attr(
|
||||
new_out_name, out_dist_attr
|
||||
)
|
||||
|
||||
# remove un-used op
|
||||
remove_op_indices = (
|
||||
group.remove_wait_op_indices
|
||||
+ group.remove_allreduce_op_indices
|
||||
+ group.remove_scale_op_indices
|
||||
)
|
||||
for idx in sorted(remove_op_indices, reverse=True):
|
||||
assert block.ops[idx].type in remove_op_types, (
|
||||
f"Unexpected: try to remove op {block.ops[idx]}"
|
||||
)
|
||||
block._remove_op(idx, False)
|
||||
|
||||
# insert coalesce op
|
||||
grad_names = [grad.name for grad in group.gradients]
|
||||
coalesce_op = block._insert_op_without_sync(
|
||||
group.coalesce_op_idx,
|
||||
type="coalesce_tensor",
|
||||
inputs={"Input": grad_names},
|
||||
outputs={
|
||||
"Output": grad_names,
|
||||
"FusedOutput": group.coalesce_var,
|
||||
},
|
||||
attrs={
|
||||
"copy_data": False,
|
||||
"use_align": True,
|
||||
"dtype": group.dtype,
|
||||
"concated_shapes": concated_shapes,
|
||||
"concated_ranks": concated_ranks,
|
||||
OP_ROLE_KEY: OpRole.Backward,
|
||||
},
|
||||
)
|
||||
|
||||
op_dist_attr = OperatorDistAttr()
|
||||
op_dist_attr.impl_idx = 0
|
||||
op_dist_attr.impl_type = "default"
|
||||
op_dist_attr.process_mesh = ProcessMesh(list(ref_process_mesh))
|
||||
for in_name in coalesce_op.input_arg_names:
|
||||
in_var = block.var(in_name)
|
||||
in_var_dist_attr = (
|
||||
self.dist_context.get_tensor_dist_attr_for_program(in_var)
|
||||
)
|
||||
op_dist_attr.set_input_dims_mapping(
|
||||
in_name, in_var_dist_attr.dims_mapping
|
||||
)
|
||||
for out_name in coalesce_op.output_arg_names:
|
||||
out_var = block.var(out_name)
|
||||
out_var_dist_attr = (
|
||||
self.dist_context.get_tensor_dist_attr_for_program(out_var)
|
||||
)
|
||||
op_dist_attr.set_output_dims_mapping(
|
||||
out_name, out_var_dist_attr.dims_mapping
|
||||
)
|
||||
|
||||
self.dist_context.set_op_dist_attr_for_program(
|
||||
coalesce_op, op_dist_attr
|
||||
)
|
||||
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _add_dependencies(self, grad_groups):
|
||||
# NOTE Currently, auto_parallel need to adopt for two executors: Sequential executor (old exe) and Graph based
|
||||
# multiple stream executor(standalone exe). This function just for standalone exe. Refactor here
|
||||
# in future when only one executor stay.
|
||||
|
||||
if len(grad_groups) == 0:
|
||||
return
|
||||
block = default_main_program().global_block()
|
||||
|
||||
# Build maps
|
||||
coalesce_to_vars_map = {}
|
||||
for group in grad_groups:
|
||||
coalesce_to_vars_map[group.coalesce_var.name] = group
|
||||
|
||||
# analyze dependencies
|
||||
dep_map = {}
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_forward_op(op):
|
||||
break
|
||||
if is_optimize_op(op):
|
||||
continue
|
||||
|
||||
if is_data_parallel_reduce_op(op):
|
||||
coalesce_var_name = op.output_arg_names[0]
|
||||
if self.coalesce_prefix in coalesce_var_name:
|
||||
group = coalesce_to_vars_map[coalesce_var_name]
|
||||
dep_map[idx] = [
|
||||
(
|
||||
idx,
|
||||
group.gradients[-1],
|
||||
group.coalesce_var,
|
||||
op.attr(OP_ROLE_KEY),
|
||||
)
|
||||
]
|
||||
dep_map[idx].append(
|
||||
(
|
||||
idx + 1,
|
||||
group.coalesce_var,
|
||||
group.gradients,
|
||||
op.attr(OP_ROLE_KEY),
|
||||
)
|
||||
)
|
||||
|
||||
# insert dependency op
|
||||
indice = sorted(dep_map.keys(), reverse=True)
|
||||
for i in indice:
|
||||
for idx, prior_vars, post_vars, op_role in dep_map[i][::-1]:
|
||||
depend_op = insert_dependencies_for_vars(
|
||||
block,
|
||||
idx,
|
||||
prior_vars,
|
||||
post_vars,
|
||||
self.dist_context,
|
||||
op_role,
|
||||
is_recompute=False,
|
||||
sync=False,
|
||||
op_namescope="data_parallel_overlap_dep",
|
||||
)
|
||||
depend_op.dist_attr.execution_stream = self.gradient_sync_stream
|
||||
block._sync_with_cpp()
|
||||
|
||||
# remove naive synchronization & assign allreduce stream
|
||||
def remove_cond(op):
|
||||
if op.type != "c_wait_compute":
|
||||
return False
|
||||
if len(op.input_arg_names) != 0:
|
||||
return False
|
||||
if len(op.output_arg_names) != 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if is_data_parallel_reduce_op(op):
|
||||
op._set_attr('use_calc_stream', True)
|
||||
op.dist_attr.execution_stream = self.gradient_sync_stream
|
||||
|
||||
if remove_cond(op):
|
||||
block._remove_op(idx, sync=False)
|
||||
|
||||
block._sync_with_cpp()
|
||||
|
||||
def summary(self, grad_groups=[]):
|
||||
# TODO: add logger module
|
||||
import logging
|
||||
|
||||
self._logger = logging.getLogger()
|
||||
self._logger.propagate = False
|
||||
if not self._logger.handlers:
|
||||
self._logger.setLevel(logging.INFO)
|
||||
log_handler = logging.StreamHandler()
|
||||
log_format = logging.Formatter(
|
||||
'[%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
|
||||
)
|
||||
log_handler.setFormatter(log_format)
|
||||
self._logger.addHandler(log_handler)
|
||||
|
||||
if len(grad_groups) > 0:
|
||||
self._logger.info("Data Parallel Optimization: ")
|
||||
self._logger.info(
|
||||
f" {len(self._grad_name_to_group_map.keys())} Allreduce ops are fused into {len(grad_groups)} coalesce allreduce ops."
|
||||
)
|
||||
self._logger.debug("gradient fusing group are following: ")
|
||||
fused_grads = set()
|
||||
for i, group in enumerate(grad_groups):
|
||||
self._logger.debug(
|
||||
f"coalesce gradient [{i}] is composed by: {[grad.name for grad in group.gradients]}"
|
||||
)
|
||||
fused_grads.update([grad.name for grad in group.gradients])
|
||||
individual_grads = set(self._grad_name_to_group_map.keys()) - set(
|
||||
fused_grads
|
||||
)
|
||||
self._logger.debug(
|
||||
f"the following [{len(individual_grads)}] gradients are not fused: "
|
||||
)
|
||||
self._logger.debug(f"individual gradient {individual_grads}")
|
||||
|
||||
|
||||
class GradientsGroup:
|
||||
def __init__(self, ops, max_group_size):
|
||||
self.max_group_size = max_group_size
|
||||
self.ops = ops
|
||||
|
||||
self.gradients = []
|
||||
self.numel = 0
|
||||
self.dtype = None
|
||||
self.ring_id = None
|
||||
self.coalesce_var = None
|
||||
self.coalesce_op_idx = -1
|
||||
self.allreduce_op_idx = -1
|
||||
self.scale_op_idx = -1
|
||||
self.remove_wait_op_indices = []
|
||||
self.remove_allreduce_op_indices = []
|
||||
self.remove_scale_op_indices = []
|
||||
|
||||
def acceptable(self, grad_var, ring_id):
|
||||
if len(self.gradients) == 0:
|
||||
return True
|
||||
if ring_id != self.ring_id:
|
||||
return False
|
||||
if get_var_numel(grad_var) + self.numel > self.max_group_size:
|
||||
return False
|
||||
if grad_var.dtype != self.dtype:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def add(self, grad_var, ring_id, i):
|
||||
self.gradients.append(grad_var)
|
||||
self.ring_id = ring_id
|
||||
self.dtype = grad_var.dtype
|
||||
self.numel += get_var_numel(grad_var)
|
||||
|
||||
# remove auxiliary ops in non-fuse dp allreduce
|
||||
self.remove_allreduce_op_indices.append(i)
|
||||
|
||||
# NOTE this pass rely on the original synchronization add in previous passes
|
||||
# (same stream or calc_wait_comm & comm_wait_calc)
|
||||
# to guarantee the correctness of comm_calc execution order.
|
||||
# so the calc_wait_comm should be keep.
|
||||
grad_op_idx = i - 1
|
||||
if i > 0 and self.ops[i - 1].type == 'c_wait_compute':
|
||||
self.remove_wait_op_indices.append(i - 1)
|
||||
grad_op_idx -= 1
|
||||
if i + 1 < len(self.ops) and is_data_parallel_scale_op(self.ops[i - 1]):
|
||||
self.remove_scale_op_indices.append(i + 1)
|
||||
|
||||
if len(self.gradients) == 1:
|
||||
# TODO Remove this is a temporary hack for Tensor Parallel. the logic
|
||||
# for find grad_op should be more general.
|
||||
if (
|
||||
self.ops[grad_op_idx].type == "all_reduce"
|
||||
and self.ops[grad_op_idx].attr("reduce_type")
|
||||
== paddle.distributed.ReduceOp.SUM
|
||||
):
|
||||
grad_op_idx -= 1
|
||||
|
||||
grad_op = self.ops[grad_op_idx]
|
||||
assert grad_var.name in grad_op.output_arg_names, (
|
||||
f"grad [{grad_var.name}] should be output of {grad_op}"
|
||||
)
|
||||
self.coalesce_op_idx = grad_op_idx
|
||||
|
||||
def finalize(self):
|
||||
self.allreduce_op_idx = self.remove_allreduce_op_indices.pop()
|
||||
if len(self.remove_wait_op_indices) > 1:
|
||||
self.remove_wait_op_indices.pop()
|
||||
if len(self.remove_scale_op_indices) > 1:
|
||||
self.scale_op_idx = self.remove_scale_op_indices.pop()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,862 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
is_optimize_op,
|
||||
is_recompute_op,
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
||||
set_var_dist_attr,
|
||||
)
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ..utils.log_utils import get_logger
|
||||
from .auto_parallel_sharding import (
|
||||
_inference_data_parallel_group_for_operator,
|
||||
_is_reshard_op,
|
||||
_skip_ops,
|
||||
is_forward_op,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
logger = get_logger(logging.INFO, "FusedLinearPromotionPass")
|
||||
|
||||
_supported_optimizer_type = [
|
||||
"adam",
|
||||
"adamax",
|
||||
"adamw",
|
||||
"decayed_adagrad",
|
||||
"momentum",
|
||||
"dgc_momentum",
|
||||
"lars_momentum",
|
||||
"merged_momentum",
|
||||
"lamb",
|
||||
"sgd",
|
||||
]
|
||||
|
||||
FUSED_LINEAR_SOURCE_PATTERNS_LIST = [
|
||||
# amp_level == 'o2' or 'o3'
|
||||
{ # only MP
|
||||
"forward": ["matmul_v2", "all_reduce", "elementwise_add"],
|
||||
"backward": ["elementwise_add_grad", "matmul_v2_grad"],
|
||||
},
|
||||
{ # MP + SP
|
||||
"forward": ["matmul_v2", "reduce_scatter", "elementwise_add"],
|
||||
"backward": [
|
||||
"elementwise_add_grad",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"all_gather",
|
||||
"matmul_v2_grad",
|
||||
"all_gather",
|
||||
],
|
||||
},
|
||||
{ # DP + MP
|
||||
"forward": ["matmul_v2", "all_reduce", "elementwise_add"],
|
||||
"backward": [
|
||||
"elementwise_add_grad",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"matmul_v2_grad",
|
||||
],
|
||||
},
|
||||
{ # DP + MP + SP
|
||||
"forward": ["matmul_v2", "reduce_scatter", "elementwise_add"],
|
||||
"backward": [
|
||||
"elementwise_add_grad",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"all_gather",
|
||||
"matmul_v2_grad",
|
||||
"all_gather",
|
||||
],
|
||||
},
|
||||
# amp_level == 'o1'
|
||||
{
|
||||
"forward": ["matmul_v2", "all_reduce", "cast", "elementwise_add"],
|
||||
"backward": ["elementwise_add_grad", "matmul_v2_grad"],
|
||||
},
|
||||
{
|
||||
"forward": ["matmul_v2", "reduce_scatter", "cast", "elementwise_add"],
|
||||
"backward": [
|
||||
"elementwise_add_grad",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"all_gather",
|
||||
"all_gather",
|
||||
"matmul_v2_grad",
|
||||
],
|
||||
},
|
||||
{
|
||||
"forward": ["matmul_v2", "all_reduce", "cast", "elementwise_add"],
|
||||
"backward": [
|
||||
"elementwise_add_grad",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"matmul_v2_grad",
|
||||
],
|
||||
},
|
||||
{
|
||||
"forward": ["matmul_v2", "reduce_scatter", "cast", "elementwise_add"],
|
||||
"backward": [
|
||||
"elementwise_add_grad",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"all_reduce",
|
||||
"scale",
|
||||
"all_gather",
|
||||
"matmul_v2_grad",
|
||||
"all_gather",
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@register_pass("auto_parallel_fused_linear_promotion")
|
||||
class FusedLinearPromotionPass(PassBase):
|
||||
"""
|
||||
Apply pre-promotion that specialized for fused_linear_pass in tensor parallelism or sequence parallelism in Auto Parallel.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("dist_context", None)
|
||||
self.set_attr("global_rank", -1)
|
||||
self.set_attr("enable_sp", False)
|
||||
self.set_attr("amp_level", "o0")
|
||||
self.set_attr("params_grads", None)
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
if (not isinstance(self.get_attr("global_rank"), int)) or self.get_attr(
|
||||
"global_rank"
|
||||
) < 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
self._dist_context = self.get_attr("dist_context")
|
||||
self._global_rank = int(self.get_attr("global_rank"))
|
||||
self._params_grads = self.get_attr("params_grads")
|
||||
self._amp_level = self.get_attr("amp_level")
|
||||
self._enable_sp = self.get_attr("enable_sp")
|
||||
self._is_amp_o1 = self._amp_level == 'o1'
|
||||
self._source_patterns = {}
|
||||
self._enable_dp, self._enable_mp = self._is_enable_dp_mp(
|
||||
self._dist_context
|
||||
)
|
||||
|
||||
pattern_offset = 4 if self._is_amp_o1 else 0
|
||||
if self._enable_sp:
|
||||
if self._enable_dp:
|
||||
self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[
|
||||
3 + pattern_offset
|
||||
]
|
||||
else:
|
||||
self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[
|
||||
1 + pattern_offset
|
||||
]
|
||||
elif self._enable_mp:
|
||||
if self._enable_dp:
|
||||
self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[
|
||||
2 + pattern_offset
|
||||
]
|
||||
else:
|
||||
self._source_patterns = FUSED_LINEAR_SOURCE_PATTERNS_LIST[
|
||||
0 + pattern_offset
|
||||
]
|
||||
else:
|
||||
logger.warning("Neither of sp and mp is enabled, skip this pass")
|
||||
return
|
||||
dp_group = None
|
||||
if self._enable_dp:
|
||||
dp_group = self._collective_data_parallel_groups(
|
||||
main_program.global_block()
|
||||
)
|
||||
|
||||
# 1. get whether the current rank is first rank in mp
|
||||
self._is_first_rank = self._is_tp_sp_first_rank(
|
||||
self._dist_context, self._global_rank
|
||||
)
|
||||
logger.debug(f"before main_program: {main_program}")
|
||||
# 2. get the forward and backward op list indexes in source patterns
|
||||
(
|
||||
forward_segments,
|
||||
backward_segments,
|
||||
) = self._get_forward_backward_op_segments(main_program)
|
||||
if len(forward_segments) == 0 or len(backward_segments) == 0:
|
||||
logger.warning(
|
||||
"No forward and backward op segments, skip this pass"
|
||||
)
|
||||
return
|
||||
# 3 transform the forward ops
|
||||
rename_var_names_map, deleted_bias_names = self._transform_forward(
|
||||
main_program,
|
||||
forward_segments,
|
||||
backward_segments,
|
||||
self._is_first_rank,
|
||||
self._enable_sp,
|
||||
self._is_amp_o1,
|
||||
)
|
||||
|
||||
# 4 transform the backward ops
|
||||
self._transform_backward(
|
||||
main_program,
|
||||
backward_segments,
|
||||
rename_var_names_map,
|
||||
self._is_first_rank,
|
||||
self._enable_sp,
|
||||
)
|
||||
|
||||
# 5. transform the optimizer ops
|
||||
self._transform_opt(
|
||||
main_program,
|
||||
deleted_bias_names,
|
||||
self._params_grads,
|
||||
self._is_first_rank,
|
||||
self._is_amp_o1,
|
||||
)
|
||||
logger.info(f"deleted_bias_names: {deleted_bias_names}")
|
||||
logger.debug(f"after main_program: {main_program}")
|
||||
|
||||
# 6. transform the startup program
|
||||
self._transform_startup_program(
|
||||
startup_program, deleted_bias_names, dp_group, self._is_first_rank
|
||||
)
|
||||
|
||||
def _is_tp_sp_first_rank(self, dist_context, rank):
|
||||
for process_mesh in dist_context.process_meshes:
|
||||
inner_mesh_shape = process_mesh.shape
|
||||
inner_mesh = (np.array(process_mesh.process_ids)).reshape(
|
||||
inner_mesh_shape
|
||||
)
|
||||
if len(inner_mesh_shape) == 1:
|
||||
return rank == min(process_mesh.process_ids)
|
||||
elif len(inner_mesh.shape) == 2:
|
||||
for id0 in range(inner_mesh_shape[0]):
|
||||
if rank == min(inner_mesh[id0, :]):
|
||||
return True
|
||||
elif len(inner_mesh.shape) == 3:
|
||||
for id0 in range(inner_mesh_shape[0]):
|
||||
for id1 in range(inner_mesh_shape[1]):
|
||||
if rank == min(inner_mesh[id0, id1, :]):
|
||||
return True
|
||||
else:
|
||||
raise ValueError("inner mesh shape is not supported")
|
||||
return False
|
||||
|
||||
def _is_enable_dp_mp(self, dist_context):
|
||||
for process_mesh in dist_context.process_meshes:
|
||||
inner_mesh_shape = process_mesh.shape
|
||||
inner_mesh = (np.array(process_mesh.process_ids)).reshape(
|
||||
inner_mesh_shape
|
||||
)
|
||||
if len(inner_mesh_shape) == 1:
|
||||
return False, inner_mesh_shape[0] > 1
|
||||
else:
|
||||
# DP * MP
|
||||
return inner_mesh_shape[-2] > 1, inner_mesh_shape[-1] > 1
|
||||
return False, False
|
||||
|
||||
def _get_forward_backward_op_segments(self, main_program):
|
||||
"""
|
||||
Get the operator segments according to the source patterns.
|
||||
"""
|
||||
|
||||
def can_match_pattern(
|
||||
ops, start_id, pattern, forward_matmul_inputs, is_backward=False
|
||||
):
|
||||
"""
|
||||
Check whether the ops in the range [start_id, start_id + len(pattern)] can match the pattern.
|
||||
If the ops is in forward pass, check it directly. However, when the ops is in backward pass,
|
||||
we need to additionally check whether the input of the last op in pattern is in forward_matmul_inputs to
|
||||
deal the case of enabling recompute.
|
||||
"""
|
||||
new_id = start_id
|
||||
if not is_backward:
|
||||
for op_name in pattern:
|
||||
if ops[new_id].type != op_name:
|
||||
return False
|
||||
new_id += 1
|
||||
forward_matmul_inputs.extend(ops[start_id].input_arg_names)
|
||||
return True
|
||||
else:
|
||||
for op_name in pattern:
|
||||
if ops[new_id].type != op_name:
|
||||
return False
|
||||
new_id += 1
|
||||
matmul_grad_input_names = ops[new_id - 1].input_arg_names
|
||||
# for refined-recompute
|
||||
if (
|
||||
matmul_grad_input_names[1] not in forward_matmul_inputs
|
||||
and matmul_grad_input_names[2] not in forward_matmul_inputs
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
global_block = main_program.global_block()
|
||||
forward_segments = []
|
||||
backward_segments = []
|
||||
ops_len = len(global_block.ops)
|
||||
|
||||
self._forward_patterns_len = len(self._source_patterns["forward"])
|
||||
self._backward_patterns_len = len(self._source_patterns["backward"])
|
||||
forward_matmul_inputs = []
|
||||
for id, op in enumerate(global_block.ops):
|
||||
if id > ops_len - self._backward_patterns_len:
|
||||
break
|
||||
if int(op.desc.attr('op_role')) == 0 or (
|
||||
is_recompute_op(op) and not op.type.endswith("_grad")
|
||||
): # forward
|
||||
if can_match_pattern(
|
||||
global_block.ops,
|
||||
id,
|
||||
self._source_patterns["forward"],
|
||||
forward_matmul_inputs,
|
||||
is_backward=False,
|
||||
):
|
||||
forward_segments.append(
|
||||
[id, id + self._forward_patterns_len]
|
||||
)
|
||||
elif int(op.desc.attr('op_role')) == 1: # backward
|
||||
if can_match_pattern(
|
||||
global_block.ops,
|
||||
id,
|
||||
self._source_patterns["backward"],
|
||||
forward_matmul_inputs,
|
||||
is_backward=True,
|
||||
):
|
||||
backward_segments.append(
|
||||
[id, id + self._backward_patterns_len]
|
||||
)
|
||||
else:
|
||||
pass
|
||||
assert len(forward_segments) >= len(backward_segments), (
|
||||
"The number of forward segments should be not shorter than the number of backward segments."
|
||||
)
|
||||
logger.info(f"forward_segments: {forward_segments}")
|
||||
logger.info(f"backward_segments: {backward_segments}")
|
||||
return forward_segments, backward_segments
|
||||
|
||||
def _collective_data_parallel_groups(self, main_block):
|
||||
for op in main_block.ops:
|
||||
if not is_forward_op(op) or op.type in _skip_ops:
|
||||
continue
|
||||
# NOTE: there aren't dist_attr in the ops which reshard insert,
|
||||
# and should be skip in sharding.
|
||||
if _is_reshard_op(op):
|
||||
continue
|
||||
group = _inference_data_parallel_group_for_operator(
|
||||
self._global_rank, op, self._dist_context
|
||||
)
|
||||
if group is not None:
|
||||
return group
|
||||
return None
|
||||
|
||||
def _transform_forward(
|
||||
self,
|
||||
main_program,
|
||||
forward_segments,
|
||||
backward_segments,
|
||||
is_first_rank,
|
||||
is_sp,
|
||||
is_amp_o1,
|
||||
):
|
||||
"""
|
||||
Transform the forward pass.
|
||||
"""
|
||||
|
||||
def _transform_forward_segment(
|
||||
global_block,
|
||||
forward_segment,
|
||||
backward_segments,
|
||||
is_first_rank,
|
||||
is_sp,
|
||||
is_amp_o1,
|
||||
):
|
||||
"""
|
||||
Transform one forward segment.
|
||||
"""
|
||||
# 1. prepare the forward_segment
|
||||
# 1.1 check whether the forward_segment is right
|
||||
origin_matmul_op = global_block.ops[forward_segment[0]]
|
||||
origin_comm_op = global_block.ops[forward_segment[0] + 1]
|
||||
origin_add_op = global_block.ops[forward_segment[1] - 1]
|
||||
origin_cast_op = (
|
||||
global_block.ops[forward_segment[1] - 2] if is_amp_o1 else None
|
||||
)
|
||||
origin_matmul_output_name = origin_matmul_op.output_arg_names[0]
|
||||
origin_comm_input_name = origin_comm_op.input_arg_names[0]
|
||||
assert origin_matmul_output_name == origin_comm_input_name, (
|
||||
f"The 0th op output name {origin_matmul_output_name} is not equal to the 1st op input name {origin_comm_input_name}"
|
||||
)
|
||||
origin_comm_output_name = origin_comm_op.output_arg_names[0]
|
||||
origin_add_input_names = origin_add_op.input_arg_names
|
||||
assert origin_comm_output_name == origin_add_input_names[0], (
|
||||
f"The 1st op output name {origin_comm_output_name} is not equal to the 2nd op input name {origin_add_input_names[0]}"
|
||||
)
|
||||
# 1.2 get the origin dist_attr
|
||||
origin_add_dist_attr = (
|
||||
self._dist_context.get_op_dist_attr_for_program(origin_add_op)
|
||||
)
|
||||
assert origin_add_dist_attr is not None, (
|
||||
f"Origin add op {origin_add_op.type} has no dist attr"
|
||||
)
|
||||
ref_mesh = origin_add_dist_attr.process_mesh
|
||||
in_var_dist_attr = origin_add_dist_attr.get_input_dist_attr(
|
||||
origin_add_op.input_arg_names[0]
|
||||
)
|
||||
ref_mapping = in_var_dist_attr.dims_mapping
|
||||
|
||||
# 2. deal matmul_v2 op
|
||||
origin_matmul_output_new_name = unique_name.generate(
|
||||
origin_matmul_output_name + "@promote"
|
||||
)
|
||||
origin_matmul_output_new_var = global_block.create_var(
|
||||
name=origin_matmul_output_new_name,
|
||||
dtype=global_block.var(origin_matmul_output_name).dtype,
|
||||
shape=global_block.var(origin_matmul_output_name).shape,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
set_var_dist_attr(
|
||||
self._dist_context,
|
||||
origin_matmul_output_new_var,
|
||||
ref_mapping,
|
||||
ref_mesh,
|
||||
)
|
||||
rename_vars_map[origin_matmul_output_name] = (
|
||||
origin_matmul_output_new_name
|
||||
)
|
||||
origin_matmul_op._rename_output(
|
||||
origin_matmul_output_name, origin_matmul_output_new_name
|
||||
)
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
origin_matmul_op, ref_mesh, ref_mapping, self._dist_context
|
||||
)
|
||||
|
||||
# 3. deal add op and cast op
|
||||
if is_first_rank:
|
||||
# insert the "elementwise_add" op before reduce_sum
|
||||
new_add_op = global_block._insert_op_without_sync(
|
||||
forward_segment[0] + 1,
|
||||
type="nop",
|
||||
)
|
||||
new_op_desc = new_add_op.desc
|
||||
new_op_desc.copy_from(origin_add_op.desc)
|
||||
# create new var of new_add_op output
|
||||
origin_add_output_name = origin_add_op.output_arg_names[0]
|
||||
new_add_op_output_name = unique_name.generate(
|
||||
origin_add_output_name + "@promote"
|
||||
)
|
||||
new_shape_var_name = (
|
||||
origin_add_output_name
|
||||
if not is_sp
|
||||
else origin_matmul_output_name
|
||||
)
|
||||
global_block.create_var(
|
||||
name=new_add_op_output_name,
|
||||
dtype=global_block.var(origin_add_output_name).dtype,
|
||||
shape=global_block.var(new_shape_var_name).shape,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
global_block._remove_var(
|
||||
origin_matmul_output_name
|
||||
) # We can remove the origin_matmul_output now.
|
||||
global_block._remove_var(origin_add_output_name)
|
||||
new_add_op._rename_output(
|
||||
origin_add_output_name, new_add_op_output_name
|
||||
)
|
||||
rename_vars_map[origin_add_op.input_arg_names[0]] = (
|
||||
origin_matmul_output_new_name
|
||||
)
|
||||
new_add_op._rename_input(
|
||||
origin_add_op.input_arg_names[0],
|
||||
origin_matmul_output_new_name,
|
||||
)
|
||||
# deal dist_attr
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
new_add_op, ref_mesh, ref_mapping, self._dist_context
|
||||
)
|
||||
# 'cast' op also need to adjust
|
||||
if is_amp_o1:
|
||||
new_cast_op = global_block._insert_op_without_sync(
|
||||
forward_segment[0] + 1,
|
||||
type="nop",
|
||||
)
|
||||
new_op_desc = new_cast_op.desc
|
||||
new_op_desc.copy_from(origin_cast_op.desc)
|
||||
if (
|
||||
new_cast_op.input_arg_names[0]
|
||||
not in delete_bias_vars_name
|
||||
): # fp16 = cast(fp32)
|
||||
delete_bias_vars_name.append(
|
||||
new_cast_op.input_arg_names[0]
|
||||
)
|
||||
else:
|
||||
if (
|
||||
new_add_op.input_arg_names[1]
|
||||
not in delete_bias_vars_name
|
||||
):
|
||||
delete_bias_vars_name.append(
|
||||
new_add_op.input_arg_names[1]
|
||||
)
|
||||
else:
|
||||
# We can remove the origin_matmul_output now.
|
||||
origin_add_output_name = origin_add_op.output_arg_names[0]
|
||||
global_block._remove_var(origin_add_output_name)
|
||||
global_block._remove_var(origin_matmul_output_name)
|
||||
|
||||
# 4. deal comm op
|
||||
# The input of all_reduce_sum only be used once, so we don't need add it in the rename_vars_map
|
||||
if is_first_rank:
|
||||
origin_comm_op._rename_input(
|
||||
origin_comm_op.input_arg_names[0],
|
||||
new_add_op.output_arg_names[0],
|
||||
)
|
||||
else:
|
||||
origin_comm_op._rename_input(
|
||||
origin_comm_op.input_arg_names[0],
|
||||
origin_matmul_output_new_name,
|
||||
)
|
||||
if (
|
||||
origin_comm_op.type == "all_reduce"
|
||||
and origin_comm_op.attr("reduce_type")
|
||||
== paddle.distributed.ReduceOp.SUM
|
||||
):
|
||||
new_comm_var_name = origin_comm_op.input_arg_names[0]
|
||||
else:
|
||||
new_comm_var_name = unique_name.generate(
|
||||
origin_comm_output_name + "@promote"
|
||||
)
|
||||
global_block.create_var(
|
||||
name=new_comm_var_name,
|
||||
dtype=global_block.var(origin_comm_output_name).dtype,
|
||||
shape=global_block.var(origin_comm_output_name).shape,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
rename_vars_map[origin_comm_output_name] = new_comm_var_name
|
||||
if global_block.has_var(origin_comm_output_name):
|
||||
global_block._remove_var(origin_comm_output_name)
|
||||
rename_vars_map[origin_add_output_name] = (
|
||||
new_comm_var_name # the output of comm op inplace the output of add op for next ops
|
||||
)
|
||||
origin_comm_op._rename_output(
|
||||
origin_comm_output_name, new_comm_var_name
|
||||
)
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
origin_comm_op, ref_mesh, ref_mapping, self._dist_context
|
||||
)
|
||||
|
||||
# 5. remove elementwise_add op and cast op
|
||||
if is_first_rank:
|
||||
if is_amp_o1:
|
||||
global_block._remove_op(forward_segment[0] + 5)
|
||||
global_block._remove_op(forward_segment[0] + 4)
|
||||
else:
|
||||
global_block._remove_op(forward_segment[0] + 3)
|
||||
else:
|
||||
global_block._remove_op(
|
||||
forward_segment[1] - 1
|
||||
) # remove elementwise_add op
|
||||
if is_amp_o1:
|
||||
if (
|
||||
origin_cast_op.input_arg_names[0]
|
||||
not in delete_bias_vars_name
|
||||
):
|
||||
delete_bias_vars_name.append(
|
||||
origin_cast_op.input_arg_names[0]
|
||||
)
|
||||
global_block._remove_var(origin_cast_op.output_arg_names[0])
|
||||
global_block._remove_op(
|
||||
forward_segment[1] - 2
|
||||
) # remove cast op
|
||||
else:
|
||||
if origin_add_input_names[1] not in delete_bias_vars_name:
|
||||
delete_bias_vars_name.append(origin_add_input_names[1])
|
||||
# update backward forward_segment
|
||||
for back_seg in reversed(backward_segments):
|
||||
if is_amp_o1:
|
||||
if back_seg[0] > forward_segment[0]:
|
||||
back_seg[0] -= 2
|
||||
back_seg[1] -= 2
|
||||
else:
|
||||
break
|
||||
else:
|
||||
if back_seg[0] > forward_segment[0]:
|
||||
back_seg[0] -= 1
|
||||
back_seg[1] -= 1
|
||||
else:
|
||||
break
|
||||
|
||||
global_block = main_program.global_block()
|
||||
rename_vars_map = {} # origin_name -> new_name
|
||||
delete_bias_vars_name = []
|
||||
for segment in reversed(forward_segments):
|
||||
_transform_forward_segment(
|
||||
global_block,
|
||||
segment,
|
||||
backward_segments,
|
||||
is_first_rank,
|
||||
is_sp,
|
||||
is_amp_o1,
|
||||
)
|
||||
global_block._sync_with_cpp()
|
||||
return rename_vars_map, delete_bias_vars_name
|
||||
|
||||
def _transform_backward(
|
||||
self,
|
||||
main_program,
|
||||
backward_segments,
|
||||
rename_var_names_map,
|
||||
is_first_rank,
|
||||
is_sp,
|
||||
):
|
||||
global_block = main_program.global_block()
|
||||
to_delete_grad_of_param = []
|
||||
if is_first_rank:
|
||||
if is_sp:
|
||||
# place the comm_op(all_gather) before the elementwise_add_grad
|
||||
for segment in reversed(backward_segments):
|
||||
add_grad_op = global_block.ops[segment[0]]
|
||||
matmul_grad_op = global_block.ops[segment[-1] - 1]
|
||||
origin_comm_op_id = segment[-1] - 2
|
||||
origin_comm_op = global_block.ops[origin_comm_op_id]
|
||||
new_comm_op = global_block._insert_op(
|
||||
segment[0],
|
||||
type="nop",
|
||||
)
|
||||
new_comm_op.desc.copy_from(origin_comm_op.desc)
|
||||
# rename input and output
|
||||
new_comm_op._rename_input(
|
||||
origin_comm_op.input_arg_names[0],
|
||||
add_grad_op.input_arg_names[0],
|
||||
)
|
||||
add_grad_op._rename_input(
|
||||
add_grad_op.input_arg_names[0],
|
||||
new_comm_op.output_arg_names[0],
|
||||
)
|
||||
matmul_grad_op._rename_input(
|
||||
matmul_grad_op.input_arg_names[0],
|
||||
add_grad_op.output_arg_names[0],
|
||||
)
|
||||
|
||||
global_block._remove_op(segment[-1] - 1)
|
||||
if self._enable_dp:
|
||||
global_block._remove_op(segment[0] + 5) # scale
|
||||
global_block._remove_op(
|
||||
segment[0] + 4
|
||||
) # all_reduce_sum
|
||||
else:
|
||||
global_block._remove_op(segment[0] + 3) # scale
|
||||
global_block._remove_op(
|
||||
segment[0] + 2
|
||||
) # all_reduce_sum
|
||||
global_block._sync_with_cpp()
|
||||
else: # not is_first_rank_in tp or sp
|
||||
# need to delete the grad op associated with the deleted bias var
|
||||
if not is_sp:
|
||||
for segment in reversed(backward_segments):
|
||||
add_grad_op = global_block.ops[segment[0]]
|
||||
rename_var_names_map[add_grad_op.output_arg_names[0]] = (
|
||||
add_grad_op.input_arg_names[0]
|
||||
)
|
||||
global_block._remove_var(add_grad_op.output_arg_names[0])
|
||||
to_delete_grad_of_param.append(
|
||||
add_grad_op.output_arg_names[1]
|
||||
)
|
||||
if self._enable_dp:
|
||||
global_block._remove_op(segment[0] + 2) # scale op
|
||||
global_block._remove_op(
|
||||
segment[0] + 1
|
||||
) # all_reduce_sum op
|
||||
global_block._remove_op(segment[0])
|
||||
global_block._sync_with_cpp()
|
||||
else:
|
||||
for segment in reversed(backward_segments):
|
||||
add_grad_op = global_block.ops[segment[0]]
|
||||
origin_comm_op = global_block.ops[segment[-1] - 2]
|
||||
rename_var_names_map[add_grad_op.output_arg_names[0]] = (
|
||||
add_grad_op.input_arg_names[0]
|
||||
)
|
||||
origin_comm_op._rename_input(
|
||||
origin_comm_op.input_arg_names[0],
|
||||
add_grad_op.input_arg_names[0],
|
||||
)
|
||||
global_block._remove_var(add_grad_op.output_arg_names[0])
|
||||
to_delete_grad_of_param.append(
|
||||
add_grad_op.output_arg_names[1]
|
||||
)
|
||||
if self._enable_dp: # DP
|
||||
global_block._remove_op(
|
||||
segment[0] + 4
|
||||
) # scale op for dp
|
||||
global_block._remove_op(
|
||||
segment[0] + 3
|
||||
) # all_reduce_sum op for dp
|
||||
global_block._remove_op(segment[0] + 2) # scale op for sp
|
||||
global_block._remove_op(
|
||||
segment[0] + 1
|
||||
) # all_reduce_sum op for sp
|
||||
global_block._remove_op(
|
||||
segment[0]
|
||||
) # elementwise_add_grad op
|
||||
global_block._sync_with_cpp()
|
||||
|
||||
# rename input vars in global_block
|
||||
for op in global_block.ops:
|
||||
if is_optimize_op(op):
|
||||
continue
|
||||
for var_name in op.input_arg_names:
|
||||
if var_name in rename_var_names_map:
|
||||
op._rename_input(var_name, rename_var_names_map[var_name])
|
||||
if self._is_amp_o1:
|
||||
for var_name in to_delete_grad_of_param:
|
||||
global_block._remove_var(var_name)
|
||||
global_block._sync_with_cpp()
|
||||
|
||||
def _transform_opt(
|
||||
self,
|
||||
main_program,
|
||||
deleted_bias_names,
|
||||
params_grads,
|
||||
is_first_rank,
|
||||
is_amp_o1,
|
||||
):
|
||||
if is_first_rank:
|
||||
return
|
||||
deleted_bias_grads_names = []
|
||||
to_delete_params_grads = []
|
||||
for id, (param, grad) in enumerate(params_grads):
|
||||
if param.name in deleted_bias_names:
|
||||
deleted_bias_grads_names.append(grad.name)
|
||||
to_delete_params_grads.append(id)
|
||||
|
||||
to_delete_op_ids = []
|
||||
for id in reversed(range(len(main_program.global_block().ops))):
|
||||
global_block = main_program.global_block()
|
||||
op = global_block.ops[id]
|
||||
op_input_names = op.input_arg_names
|
||||
for op_input in op_input_names:
|
||||
if op_input in deleted_bias_grads_names:
|
||||
if op.type in _supported_optimizer_type:
|
||||
for output_var in op.output_arg_names:
|
||||
global_block._remove_var(output_var)
|
||||
grad_var = op.input('Grad')[0]
|
||||
global_block._remove_var(grad_var)
|
||||
to_delete_op_ids.append(id)
|
||||
if (
|
||||
op.type == "squared_l2_norm"
|
||||
or op.type == "clip_by_norm"
|
||||
):
|
||||
output_var_name = op.output_arg_names[0]
|
||||
global_block._remove_var(output_var_name)
|
||||
to_delete_op_ids.append(id)
|
||||
for intra_id in range(id + 1, len(global_block.ops)):
|
||||
intra_op = global_block.ops[intra_id]
|
||||
if (
|
||||
output_var_name in intra_op.input_arg_names
|
||||
and intra_op.type == "stack"
|
||||
):
|
||||
origin_vars = intra_op.input("X")
|
||||
origin_vars.remove(output_var_name)
|
||||
intra_op.desc.set_input("X", origin_vars)
|
||||
break
|
||||
if op.type == "elementwise_mul":
|
||||
to_delete_op_ids.append(id)
|
||||
# check_finite_and_unscale and update_loss_scaling
|
||||
if (
|
||||
op.type == "check_finite_and_unscale"
|
||||
or op.type == "update_loss_scaling"
|
||||
):
|
||||
origin_vars = op.input("X")
|
||||
origin_vars.remove(op_input)
|
||||
op.desc.set_input("X", origin_vars)
|
||||
origin_vars = op.output("Out")
|
||||
origin_vars.remove(op_input)
|
||||
op.desc.set_output("Out", origin_vars)
|
||||
|
||||
if is_amp_o1:
|
||||
for output_name in op.output_arg_names:
|
||||
if (
|
||||
output_name in deleted_bias_grads_names
|
||||
and op.type == 'cast'
|
||||
):
|
||||
to_delete_op_ids.append(id)
|
||||
|
||||
for id in to_delete_op_ids:
|
||||
global_block._remove_op(id)
|
||||
main_program.global_block()._sync_with_cpp()
|
||||
|
||||
for id in reversed(to_delete_params_grads):
|
||||
del params_grads[id]
|
||||
return
|
||||
|
||||
def _transform_startup_program(
|
||||
self, startup_program, deleted_bias_names, dp_group, is_first_rank
|
||||
):
|
||||
"""
|
||||
Delete the vars and ops associated with deleted_bias_names in startup program.
|
||||
"""
|
||||
logger.debug(f"Before transform startup_program: {startup_program}")
|
||||
cur_glock = startup_program.global_block()
|
||||
to_delete_op_ids = []
|
||||
# for variables associated with deleted_bias_names in amp-o2, such as 'opt_linear_1.b_0_fp32_master_0'
|
||||
to_delete_extra_vars = []
|
||||
for id, op in enumerate(cur_glock.ops):
|
||||
if not is_first_rank:
|
||||
output_var = op.output_arg_names[0]
|
||||
if output_var in deleted_bias_names:
|
||||
to_delete_op_ids.append(id)
|
||||
else:
|
||||
for var_name in deleted_bias_names:
|
||||
if var_name in output_var:
|
||||
to_delete_op_ids.append(id)
|
||||
if output_var not in to_delete_extra_vars:
|
||||
to_delete_extra_vars.append(output_var)
|
||||
else:
|
||||
if op.type == "broadcast":
|
||||
input_vars = op.input_arg_names
|
||||
if (
|
||||
input_vars[0] in deleted_bias_names
|
||||
and id not in to_delete_op_ids
|
||||
):
|
||||
if dp_group is None or (
|
||||
dp_group is not None
|
||||
and op.attr("ring_id") != dp_group.id
|
||||
):
|
||||
to_delete_op_ids.append(id)
|
||||
for to_delete_id in reversed(to_delete_op_ids):
|
||||
cur_glock._remove_op(to_delete_id)
|
||||
if not is_first_rank:
|
||||
for var_name in deleted_bias_names:
|
||||
cur_glock._remove_var(var_name)
|
||||
for var_name in to_delete_extra_vars:
|
||||
if cur_glock.has_var(var_name):
|
||||
cur_glock._remove_var(var_name)
|
||||
cur_glock._sync_with_cpp()
|
||||
logger.debug(f"After transform startup_program: {startup_program}")
|
||||
@@ -0,0 +1,538 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
from functools import reduce
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
|
||||
from ..auto_parallel.process_mesh import ProcessMesh
|
||||
from ..auto_parallel.static.dist_attribute import (
|
||||
OperatorDistAttr,
|
||||
TensorDistAttr,
|
||||
)
|
||||
from ..auto_parallel.static.operators.common import (
|
||||
SyncMode,
|
||||
is_data_parallel_reduce_op,
|
||||
)
|
||||
from ..auto_parallel.static.process_group import (
|
||||
get_all_process_groups,
|
||||
get_world_process_group,
|
||||
)
|
||||
from ..auto_parallel.static.reshard import Resharder
|
||||
from ..auto_parallel.static.utils import (
|
||||
_get_comm_group,
|
||||
insert_dependencies_for_vars,
|
||||
is_gradient_clip_op,
|
||||
is_optimize_op,
|
||||
is_reshard_op,
|
||||
)
|
||||
from .auto_parallel_sharding import ShardingPass
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
|
||||
def _get_params_grads(block):
|
||||
params_grads = []
|
||||
for op in reversed(block.ops):
|
||||
if not is_optimize_op(op):
|
||||
break
|
||||
if "Param" in op.input_names and "Grad" in op.input_names:
|
||||
param_name = op.input("Param")[0]
|
||||
grad_name = op.input("Grad")[0]
|
||||
param = block.var(param_name)
|
||||
grad = block.var(grad_name)
|
||||
params_grads.append((param, grad))
|
||||
return params_grads
|
||||
|
||||
|
||||
def _get_dpmp_topology(origin_topology, sharding_group):
|
||||
"""
|
||||
Get dpmp topology from origin_topology
|
||||
|
||||
Example:
|
||||
the parallel strategy: dp4-mp2-sharding2
|
||||
the complete process_mesh:
|
||||
topology: [4, 2]
|
||||
processes: [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
the dpmp topology: [2, 2]
|
||||
the sharding axis: 1
|
||||
"""
|
||||
sharding_axis = 1
|
||||
dp_sharding_topology = [
|
||||
origin_topology[0] // sharding_group.nranks,
|
||||
sharding_group.nranks,
|
||||
]
|
||||
if dp_sharding_topology[0] == 1:
|
||||
sharding_axis = 0
|
||||
dp_sharding_topology = dp_sharding_topology[1:]
|
||||
|
||||
product_dp_sharding = reduce(lambda x, y: x * y, dp_sharding_topology, 1)
|
||||
product_topology = reduce(lambda x, y: x * y, origin_topology, 1)
|
||||
|
||||
if product_topology == product_dp_sharding:
|
||||
dpmp_topology = dp_sharding_topology
|
||||
else:
|
||||
assert product_topology % product_dp_sharding == 0
|
||||
mp_degree = product_topology // product_dp_sharding
|
||||
dpmp_topology = [*dp_sharding_topology, mp_degree]
|
||||
|
||||
return dpmp_topology, sharding_axis
|
||||
|
||||
|
||||
def _get_dpmp_process_mesh(rank_id, topology, processes, sharding_group):
|
||||
"""
|
||||
Get dpmp process_mesh from the complete process_mesh which apply sharding.
|
||||
|
||||
Example:
|
||||
the parallel strategy: dp4-mp2-sharding2
|
||||
the complete process_mesh:
|
||||
topology: [4, 2]
|
||||
processes: [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
the dpmp process_mesh is:
|
||||
1) topology: [2, 2], processes: [0, 1, 4, 5]
|
||||
2) topology: [2, 2], processes: [2, 3, 6, 7]
|
||||
"""
|
||||
if sharding_group is None:
|
||||
return topology, processes
|
||||
|
||||
# get dpmp_topology
|
||||
dpmp_topology, sharding_axis = _get_dpmp_topology(topology, sharding_group)
|
||||
|
||||
# get all sharding_groups of ranks
|
||||
sharding_groups = []
|
||||
for rank in processes:
|
||||
group = _get_comm_group(processes, dpmp_topology, sharding_axis, rank)
|
||||
if group not in sharding_groups:
|
||||
sharding_groups.append(group)
|
||||
|
||||
# get dpmp_processes
|
||||
sharding_groups = np.array(sharding_groups)
|
||||
dpmp_processes_in_sharding = None
|
||||
for i in range(sharding_groups.shape[-1]):
|
||||
if rank_id in sharding_groups[:, i]:
|
||||
dpmp_processes_in_sharding = sharding_groups[:, i]
|
||||
|
||||
assert dpmp_processes_in_sharding is not None
|
||||
return dpmp_topology, list(dpmp_processes_in_sharding)
|
||||
|
||||
|
||||
def _is_about_global_norm(
|
||||
rank_id, tensor_shape, topology, processes, dims_mapping, sharding_group
|
||||
):
|
||||
# get current process_mesh where the parameter exist.
|
||||
dpmp_topology, dpmp_processes = _get_dpmp_process_mesh(
|
||||
rank_id, topology, processes, sharding_group
|
||||
)
|
||||
|
||||
complete_shape = Resharder.compute_complete_shape(
|
||||
tensor_shape, dpmp_topology, dims_mapping
|
||||
)
|
||||
|
||||
complete_partitions = []
|
||||
complete_param_ranks = []
|
||||
for process in dpmp_processes:
|
||||
partition_index = Resharder.compute_partition_index(
|
||||
process, complete_shape, dims_mapping, dpmp_topology, dpmp_processes
|
||||
)
|
||||
if partition_index not in complete_partitions:
|
||||
complete_partitions.append(partition_index)
|
||||
complete_param_ranks.append(process)
|
||||
|
||||
return rank_id in complete_param_ranks
|
||||
|
||||
|
||||
class ClipHelper:
|
||||
def __init__(
|
||||
self, params_grads, rank_id, block, dist_context, pass_context
|
||||
):
|
||||
params, _ = zip(*params_grads)
|
||||
self.params = list(params)
|
||||
self.params_name = [p.name for p in self.params]
|
||||
self.rank_id = rank_id
|
||||
self.block = block
|
||||
self.dist_context = dist_context
|
||||
self.pass_context = pass_context
|
||||
self.sharding_group = None
|
||||
self.world_ranks = get_world_process_group().ranks
|
||||
if hasattr(dist_context, '_sharding_group'):
|
||||
self.sharding_group = dist_context._sharding_group
|
||||
|
||||
self.world_nranks = len(self.world_ranks)
|
||||
self.pure_data_parallel = self._is_pure_data_parallel()
|
||||
self.rank_to_params = self._partition_parameters(params)
|
||||
|
||||
def is_calculate_norm(self, name):
|
||||
"""
|
||||
whether the param_name@GRAD participate in the calculation of global_norm
|
||||
"""
|
||||
if not self.is_local_param(name):
|
||||
return False
|
||||
|
||||
param = self.params[self.params_name.index(name)]
|
||||
if not self.pure_data_parallel:
|
||||
dist_attr = self._get_dist_attr(name)
|
||||
topology = dist_attr.process_mesh.shape
|
||||
processes = dist_attr.process_mesh.process_ids
|
||||
dims_mapping = dist_attr.dims_mapping
|
||||
return _is_about_global_norm(
|
||||
self.rank_id,
|
||||
param.shape,
|
||||
topology,
|
||||
processes,
|
||||
dims_mapping,
|
||||
self.sharding_group,
|
||||
)
|
||||
else:
|
||||
return param.name in self.rank_to_params[self.rank_id]
|
||||
|
||||
def is_local_param(self, name):
|
||||
"""
|
||||
whether the param_name is updated with opt in cur_rank
|
||||
"""
|
||||
if name not in self.params_name:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _get_dist_attr(self, name):
|
||||
var = self.block.vars[name]
|
||||
return self.dist_context.get_tensor_dist_attr_for_program(var)
|
||||
|
||||
def is_local_var_with_dist_attr(self, name):
|
||||
"""
|
||||
whether the var_name is belong to cur_rank
|
||||
"""
|
||||
dist_attr = self._get_dist_attr(name)
|
||||
assert dist_attr is not None
|
||||
return self.rank_id in dist_attr.process_mesh.process_ids
|
||||
|
||||
def _init_dist_attr(self, op):
|
||||
op_dist_attr = OperatorDistAttr()
|
||||
op_dist_attr.process_mesh = ProcessMesh(self.world_ranks)
|
||||
for in_name in op.input_arg_names:
|
||||
in_var = self.block.vars[in_name]
|
||||
in_dist_attr = TensorDistAttr()
|
||||
in_dist_attr.process_mesh = ProcessMesh(self.world_ranks)
|
||||
in_dist_attr.dims_mapping = [-1 for i in in_var.shape]
|
||||
self.dist_context.set_tensor_dist_attr_for_program(
|
||||
in_var, in_dist_attr
|
||||
)
|
||||
op_dist_attr.set_input_dist_attr(in_name, in_dist_attr)
|
||||
for out_name in op.output_arg_names:
|
||||
out_var = self.block.vars[out_name]
|
||||
out_dist_attr = TensorDistAttr()
|
||||
out_dist_attr.process_mesh = ProcessMesh(self.world_ranks)
|
||||
out_dist_attr.dims_mapping = [-1 for i in out_var.shape]
|
||||
self.dist_context.set_tensor_dist_attr_for_program(
|
||||
out_var, out_dist_attr
|
||||
)
|
||||
op_dist_attr.set_output_dist_attr(out_name, out_dist_attr)
|
||||
self.dist_context.set_op_dist_attr_for_program(op, op_dist_attr)
|
||||
|
||||
def _is_pure_data_parallel(self):
|
||||
for applied_pass in self.pass_context.passes:
|
||||
if isinstance(applied_pass, ShardingPass):
|
||||
return False
|
||||
|
||||
groups = get_all_process_groups()
|
||||
for g in groups:
|
||||
if g.nranks != self.world_nranks:
|
||||
return False
|
||||
|
||||
for op in self.block.ops:
|
||||
if (
|
||||
(
|
||||
op.type == "reduce"
|
||||
and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
|
||||
)
|
||||
or (
|
||||
op.type == "all_reduce"
|
||||
and op.desc.attr("reduce_type") == dist.ReduceOp.SUM
|
||||
)
|
||||
and not is_data_parallel_reduce_op(op)
|
||||
):
|
||||
return False
|
||||
if op.type in ["send_v2", "recv_v2"]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _partition_parameters(self, params):
|
||||
"""
|
||||
build rank_id_to_params by the param's numel
|
||||
to guarantee params in every rank of dp_group as even as possible.
|
||||
"""
|
||||
mapping = {}
|
||||
if not self.pure_data_parallel:
|
||||
for rank_ in range(self.world_nranks):
|
||||
mapping[rank_] = [p.name for p in params]
|
||||
else:
|
||||
for rank_ in range(self.world_nranks):
|
||||
mapping[rank_] = []
|
||||
sizes = [0] * self.world_nranks
|
||||
for param in params:
|
||||
rank = sizes.index(min(sizes))
|
||||
mapping[rank].append(param.name)
|
||||
numel = reduce(lambda x, y: x * y, param.shape, 1)
|
||||
assert numel > 0, (
|
||||
f"param [{param.name}] should larger than 0, but it is [{numel}]"
|
||||
)
|
||||
sizes[rank] += numel
|
||||
return mapping
|
||||
|
||||
|
||||
@register_pass("auto_parallel_grad_clip")
|
||||
class ClipGradByGlobalNormPass(PassBase):
|
||||
"""
|
||||
1. Remove norm-compute op and grad-scale op when the grad is not in current rank
|
||||
or is independent of the calculation of norm.
|
||||
2. Each rank computes its own norm value, then gets global_norm by allreduce_sum only once.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("rank_id", None)
|
||||
self.set_attr("dist_context", None)
|
||||
self.set_attr("params_grads", None)
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
dist_context = self.get_attr("dist_context")
|
||||
if dist_context._serial_optimizer._grad_clip is None:
|
||||
return False
|
||||
if self.get_attr("params_grads") is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
dist_context = self.get_attr("dist_context", None)
|
||||
rank_id = self.get_attr("rank_id", None)
|
||||
block = main_program.global_block()
|
||||
dist_params_grads = self.get_attr("params_grads", None)
|
||||
# dist_params_grads = _get_params_grads(block)
|
||||
|
||||
self.clip_helper = ClipHelper(
|
||||
dist_params_grads, rank_id, block, dist_context, context
|
||||
)
|
||||
self._remove_no_need_ops_vars(block)
|
||||
|
||||
def _remove_no_need_ops_vars(self, block):
|
||||
removed_op_out_type = [
|
||||
'squared_l2_norm',
|
||||
'square',
|
||||
'reduce_sum',
|
||||
]
|
||||
|
||||
removed_op_idx = set()
|
||||
removed_tmp_var = set()
|
||||
for idx, op in enumerate(block.ops):
|
||||
if not is_gradient_clip_op(op):
|
||||
continue
|
||||
|
||||
if op.type == 'clip_by_norm':
|
||||
# remove 'clip_by_norm' op if the param is not updated with opt in current rank
|
||||
input_name = op.input("X")[0]
|
||||
if input_name.find("@GRAD") != -1:
|
||||
param_name = input_name[: input_name.find("@GRAD")]
|
||||
is_local = self.clip_helper.is_local_param(param_name)
|
||||
if not is_local:
|
||||
removed_op_idx.add(idx)
|
||||
removed_tmp_var.update(set(op.output_arg_names))
|
||||
|
||||
elif op.type in removed_op_out_type:
|
||||
input_name = op.input("X")[0]
|
||||
if input_name.find("@GRAD") != -1:
|
||||
# remove 'squared_l2_norm' and 'square' ops,
|
||||
# if the param@GRAD in cur_rank does not participate in the calculation of global_norm
|
||||
param_name = input_name[: input_name.find("@GRAD")]
|
||||
is_local = self.clip_helper.is_local_param(param_name)
|
||||
is_calculate = self.clip_helper.is_calculate_norm(
|
||||
param_name
|
||||
)
|
||||
if not is_local or not is_calculate:
|
||||
removed_op_idx.add(idx)
|
||||
removed_tmp_var.update(set(op.output_arg_names))
|
||||
else:
|
||||
# 'reduce_sum' must be behind 'square'
|
||||
if idx - 1 in removed_op_idx:
|
||||
removed_op_idx.add(idx)
|
||||
removed_tmp_var.update(set(op.output_arg_names))
|
||||
|
||||
elif op.type == 'elementwise_mul':
|
||||
# 'elementwise_mul' scale the param@GRAD with global_norm
|
||||
# remove 'elementwise_mul' op if the param is not updated with opt in current rank
|
||||
input_name = op.input("X")[0]
|
||||
if input_name.find("@GRAD") != -1:
|
||||
param_name = input_name[: input_name.find("@GRAD")]
|
||||
is_local = self.clip_helper.is_local_param(param_name)
|
||||
if not is_local:
|
||||
removed_op_idx.add(idx)
|
||||
if block.ops[idx - 1].type == 'cast':
|
||||
removed_op_idx.add(idx - 1)
|
||||
removed_tmp_var.update(
|
||||
set(block.ops[idx - 1].output_arg_names)
|
||||
)
|
||||
|
||||
elif op.type == 'sum':
|
||||
# 'sum' op is used to calculate global_norm, and need to filter inputs which is not in cur_rank
|
||||
reserved_vars = []
|
||||
for input_name in op.input_arg_names:
|
||||
if (
|
||||
input_name not in removed_tmp_var
|
||||
and self.clip_helper.is_local_var_with_dist_attr(
|
||||
input_name
|
||||
)
|
||||
):
|
||||
reserved_vars.append(input_name)
|
||||
if not reserved_vars:
|
||||
removed_op_idx.add(idx)
|
||||
removed_tmp_var.update(set(op.output_arg_names))
|
||||
if block.ops[idx + 1].type == 'cast':
|
||||
removed_op_idx.add(idx + 1)
|
||||
removed_tmp_var.update(
|
||||
set(block.ops[idx + 1].output_arg_names)
|
||||
)
|
||||
else:
|
||||
op.desc.set_input("X", reserved_vars)
|
||||
|
||||
elif op.type == 'stack':
|
||||
# 'stack' op is also used to calculate global_norm ('stack' + 'reduce_sum'), and need to filter inputs which is not in cur_rank
|
||||
reserved_vars = []
|
||||
for input_name in op.input_arg_names:
|
||||
if (
|
||||
input_name not in removed_tmp_var
|
||||
and self.clip_helper.is_local_var_with_dist_attr(
|
||||
input_name
|
||||
)
|
||||
):
|
||||
reserved_vars.append(input_name)
|
||||
if not reserved_vars:
|
||||
removed_op_idx.add(idx)
|
||||
removed_tmp_var.update(set(op.output_arg_names))
|
||||
if block.ops[idx + 1].type == 'reduce_sum':
|
||||
removed_op_idx.add(idx + 1)
|
||||
removed_tmp_var.update(
|
||||
set(block.ops[idx + 1].output_arg_names)
|
||||
)
|
||||
if block.ops[idx + 2].type == 'cast':
|
||||
removed_op_idx.add(idx + 2)
|
||||
removed_tmp_var.update(
|
||||
set(block.ops[idx + 2].output_arg_names)
|
||||
)
|
||||
else:
|
||||
op.desc.set_input("X", reserved_vars)
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if not (is_optimize_op(op) or is_reshard_op(op)):
|
||||
break
|
||||
if not is_gradient_clip_op(op):
|
||||
continue
|
||||
if idx in removed_op_idx:
|
||||
block._remove_op(idx, sync=False)
|
||||
|
||||
for idx, op in reversed(list(enumerate(block.ops))):
|
||||
if not (is_optimize_op(op) or is_reshard_op(op)):
|
||||
break
|
||||
if not is_gradient_clip_op(op):
|
||||
continue
|
||||
if op.type == 'sqrt':
|
||||
input_name = op.input("X")[0]
|
||||
input_var = block.vars[input_name]
|
||||
insert_leaf_fill_constant_node = False
|
||||
if paddle.distributed.get_world_size() > 1:
|
||||
offset = 0
|
||||
if input_name in removed_tmp_var:
|
||||
removed_tmp_var.remove(input_name)
|
||||
fill_constant_op = block._insert_op(
|
||||
idx,
|
||||
type='fill_constant',
|
||||
inputs={},
|
||||
outputs={'Out': [input_var]},
|
||||
attrs={
|
||||
'shape': [],
|
||||
'dtype': input_var.dtype,
|
||||
'value': 0,
|
||||
'force_cpu': False,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
fill_constant_op._set_attr(
|
||||
'op_namescope', "/gradient_clip_pass"
|
||||
)
|
||||
offset += 1
|
||||
self.clip_helper._init_dist_attr(fill_constant_op)
|
||||
insert_leaf_fill_constant_node = True
|
||||
|
||||
allreduce_op = block._insert_op(
|
||||
idx + offset,
|
||||
type='all_reduce',
|
||||
inputs={'x': [input_var]},
|
||||
outputs={'out': [input_var]},
|
||||
attrs={
|
||||
'ring_id': 0,
|
||||
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
||||
OP_ROLE_KEY: OpRole.Optimize,
|
||||
},
|
||||
)
|
||||
# TODO better regular the usage of op namescope
|
||||
allreduce_op._set_attr(
|
||||
'op_namescope', '/' + SyncMode.GlobalNormSync
|
||||
)
|
||||
self.clip_helper._init_dist_attr(allreduce_op)
|
||||
|
||||
if insert_leaf_fill_constant_node:
|
||||
# NOTE add naive deps for global norm sync in graph exe
|
||||
j = idx - 1
|
||||
prior_op = None
|
||||
while j > 0:
|
||||
op_type = block.ops[j].type
|
||||
if op_type in [
|
||||
'update_loss_scaling',
|
||||
'check_finite_and_unscale',
|
||||
] or op_type.endswith("_grad"):
|
||||
prior_op = block.ops[j]
|
||||
break
|
||||
j -= 1
|
||||
assert prior_op is not None, (
|
||||
"Unexpected: ClipByGlobalNorm could not find priory depend op"
|
||||
)
|
||||
prior_var = block.vars[prior_op.output_arg_names[0]]
|
||||
assert prior_var is not None, (
|
||||
"Unexpected: ClipByGlobalNorm could not find priory depend var"
|
||||
)
|
||||
insert_dependencies_for_vars(
|
||||
block,
|
||||
idx,
|
||||
prior_var,
|
||||
input_var,
|
||||
self.clip_helper.dist_context,
|
||||
OpRole.Optimize,
|
||||
process_mesh=[
|
||||
-1
|
||||
], # hack to avoid initialize the dist attr for coalesce var
|
||||
is_recompute=False,
|
||||
sync=False,
|
||||
op_namescope="grad_clip_fill_constant_dep",
|
||||
)
|
||||
|
||||
for varname in removed_tmp_var:
|
||||
block._remove_var(varname, sync=False)
|
||||
|
||||
block._sync_with_cpp()
|
||||
@@ -0,0 +1,346 @@
|
||||
# Copyright (c) 2021 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.
|
||||
from __future__ import annotations
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.process_group import (
|
||||
get_world_process_group,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import (
|
||||
OpRole,
|
||||
)
|
||||
from paddle.framework import (
|
||||
_current_expected_place_ as _get_device,
|
||||
)
|
||||
|
||||
from .pass_base import PassBase, PassType, register_pass
|
||||
|
||||
world_process_group = get_world_process_group()
|
||||
|
||||
|
||||
def _move_used_grad_op(used_grad_op, grad):
|
||||
move_to_opt_block_flag = True
|
||||
move_to_opt_ops = []
|
||||
cannot_move_op = ["pd_op.send_v2", "pd_op.send"]
|
||||
|
||||
def find_move_op(backward_op):
|
||||
nonlocal move_to_opt_block_flag
|
||||
if not move_to_opt_block_flag or backward_op in move_to_opt_ops:
|
||||
return
|
||||
if backward_op.name() in cannot_move_op:
|
||||
move_to_opt_block_flag = False
|
||||
return
|
||||
if backward_op.num_operands() == 1:
|
||||
move_to_opt_block_flag = True
|
||||
move_to_opt_ops.append(backward_op)
|
||||
elif backward_op.name() == "pd_op.slice":
|
||||
move_to_opt_ops.append(backward_op)
|
||||
for i in range(0, backward_op.num_operands()):
|
||||
if not grad.is_same(backward_op.operand_source(i)):
|
||||
move_to_opt_ops.append(
|
||||
backward_op.operand_source(i).get_defining_op()
|
||||
)
|
||||
move_to_opt_block_flag = True
|
||||
else:
|
||||
# NOTE(zhangwl):temp only consider one operand op
|
||||
move_to_opt_block_flag = False
|
||||
return
|
||||
for op_result in backward_op.results():
|
||||
for next_op in op_result.all_used_ops():
|
||||
if next_op.op_role != int(OpRole.Optimize):
|
||||
find_move_op(next_op)
|
||||
|
||||
find_move_op(used_grad_op)
|
||||
if move_to_opt_block_flag:
|
||||
for move_op in move_to_opt_ops:
|
||||
move_op.op_role = int(OpRole.Optimize)
|
||||
|
||||
|
||||
def _pir_append_gradient_merge_backward_op(
|
||||
main_program,
|
||||
startup_program,
|
||||
params_grads,
|
||||
):
|
||||
main_block = main_program.global_block()
|
||||
startup_block = startup_program.global_block()
|
||||
|
||||
# {param: gradient_merge_var} to insert scale op and fill_constant op
|
||||
new_params_grads = []
|
||||
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)
|
||||
|
||||
for param, grad in params_grads:
|
||||
if grad is None:
|
||||
continue
|
||||
|
||||
assert not param.is_selected_row_type(), (
|
||||
"SELECTED_ROWS is not supported in GradientMergeOptimizer for now"
|
||||
)
|
||||
|
||||
grad_dtype = grad.dtype
|
||||
grad_type = grad.type()
|
||||
|
||||
for op in grad.all_used_ops():
|
||||
if op.has_attr("master_grad_cast"):
|
||||
grad_dtype = op.result(0).dtype
|
||||
grad_type = op.result(0).type()
|
||||
|
||||
# step1: create gradient_merge var and init with 0
|
||||
# Add persistable gradient variables in startup_program
|
||||
paddle.pir.set_insertion_point_to_block_end(startup_block)
|
||||
gradient_merge_var = paddle.full(
|
||||
shape=grad._local_shape, fill_value=0.0, dtype=grad_dtype
|
||||
)
|
||||
gradient_merge_var.persistable = True
|
||||
paddle.pir.set_insertion_point_after(
|
||||
gradient_merge_var.get_defining_op()
|
||||
)
|
||||
paddle._C_ops.set_persistable_value(
|
||||
gradient_merge_var, param.name + "@GRAD@MERGE"
|
||||
)
|
||||
|
||||
# step2: Accumulate persistable gradient variables in main_program
|
||||
# NOTE(zhaoyingli): inplace operation must be 'a = a + b', cannot be 'a = b + a'
|
||||
grad_defining_op = grad.get_defining_op()
|
||||
paddle.pir.set_insertion_point_after(grad_defining_op)
|
||||
|
||||
new_gradient_merge_var = main_block.add_kwarg(
|
||||
param.name + "@GRAD@MERGE", grad_type
|
||||
)
|
||||
new_gradient_merge_var.persistable = True
|
||||
new_gradient_merge_var.place_attr = cur_place
|
||||
new_gradient_merge_var_add = paddle._C_ops.add_(
|
||||
new_gradient_merge_var, grad
|
||||
)
|
||||
new_gradient_merge_var_add_op = (
|
||||
new_gradient_merge_var_add.get_defining_op()
|
||||
)
|
||||
new_gradient_merge_var_add_op.op_role = grad_defining_op.op_role
|
||||
|
||||
new_gradient_merge_var_add_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
grad_defining_op.dist_attr.process_mesh,
|
||||
grad_defining_op.dist_attr.operands(),
|
||||
grad_defining_op.dist_attr.results(),
|
||||
grad_defining_op.dist_attr.chunk_id,
|
||||
)
|
||||
)
|
||||
new_gradient_merge_var_add_op.set_bool_attr("grad_merge_add", True)
|
||||
|
||||
# NOTE(zhangweilong): grad may in different device in auto_parallel, so need consider all_gather/all_reduce/split/... op
|
||||
for used_grad_op in grad.all_used_ops():
|
||||
_move_used_grad_op(used_grad_op, grad)
|
||||
|
||||
opt_ops_use_grad = [
|
||||
op
|
||||
for op in grad.all_used_ops()
|
||||
if op.op_role == int(OpRole.Optimize)
|
||||
]
|
||||
|
||||
grad.replace_grad_users_with(
|
||||
new_gradient_merge_var, set(opt_ops_use_grad)
|
||||
)
|
||||
|
||||
# reset gradient merge var to zero after finishing optimization
|
||||
paddle.pir.set_insertion_point_to_block_end(main_block)
|
||||
set_value = paddle.full(
|
||||
shape=[1], fill_value=float(0), dtype=grad_dtype
|
||||
)
|
||||
new_gradient_merge_var_zero = paddle._C_ops.set_value_with_tensor_(
|
||||
new_gradient_merge_var, set_value, [], [], [], [], [], []
|
||||
)
|
||||
|
||||
set_value_op = new_gradient_merge_var_zero.get_defining_op()
|
||||
set_value_op.op_role = int(OpRole.Optimize)
|
||||
for id in range(1, set_value_op.num_operands()):
|
||||
op_input = set_value_op.operand_source(id)
|
||||
op_input.get_defining_op().op_role = int(OpRole.Optimize)
|
||||
|
||||
# step3: Construct new_params_grads and grad_to_gradient_merge
|
||||
new_params_grads.append((param, new_gradient_merge_var))
|
||||
|
||||
return new_params_grads
|
||||
|
||||
|
||||
def _pir_move_reduce_to_backward_stage(main_program):
|
||||
pass
|
||||
|
||||
|
||||
def _pir_remove_cast_for_master_grad(main_program, params_grads):
|
||||
for op in main_program.global_block().ops:
|
||||
if op.has_attr("master_grad_cast"):
|
||||
op.result(0).replace_all_uses_with(op.operand_source(0))
|
||||
op.erase()
|
||||
|
||||
|
||||
def _find_trivial_optimizer_ops(block):
|
||||
optimizer_ops = []
|
||||
for op in block.ops:
|
||||
if "adam" in op.name() or "sgd" in op.name():
|
||||
optimizer_ops.append(op)
|
||||
return optimizer_ops
|
||||
|
||||
|
||||
def _get_prev_op(block, optimizer_op):
|
||||
found = False
|
||||
for op in reversed(block.ops):
|
||||
if found:
|
||||
return op
|
||||
if op.id == optimizer_op.id:
|
||||
found = True
|
||||
return None
|
||||
|
||||
|
||||
def _insert_scale_op_after(target_value, optimizer_op, scale, bias=0.0):
|
||||
scaled_grad = paddle._C_ops.scale_(target_value, scale, bias, False)
|
||||
|
||||
scale_op = scaled_grad.get_defining_op()
|
||||
scale_op.op_role = int(OpRole.Optimize)
|
||||
|
||||
full_op = scale_op.operand_source(1).get_defining_op()
|
||||
assert full_op.name() == "pd_op.full", (
|
||||
f"The defining op of the scale value should be `pd_op.full`, but got {full_op.name()}"
|
||||
)
|
||||
full_op.op_role = int(OpRole.Optimize)
|
||||
|
||||
if "adam" in optimizer_op.name():
|
||||
optimizer_op.operand(1).set_source(scaled_grad)
|
||||
elif "sgd" in optimizer_op.name():
|
||||
optimizer_op.operand(2).set_source(scaled_grad)
|
||||
|
||||
|
||||
def _append_scale_op_before_comm(block, new_params_to_grads, k_steps):
|
||||
for op in reversed(block.ops):
|
||||
if op.op_role == int(OpRole.Backward):
|
||||
paddle.pir.set_insertion_point_after(op)
|
||||
break
|
||||
for _, new_grad in new_params_to_grads:
|
||||
new_grad = paddle._C_ops.scale_(new_grad, 1.0 / k_steps, 0.0, False)
|
||||
|
||||
scale_op = new_grad.get_defining_op()
|
||||
scale_op.op_role = int(OpRole.Optimize)
|
||||
|
||||
full_op = scale_op.operand_source(1).get_defining_op()
|
||||
assert full_op.name() == "pd_op.full", (
|
||||
f"The defining op of the scale value should be `pd_op.full`, but got {full_op.name()}"
|
||||
)
|
||||
full_op.op_role = int(OpRole.Optimize)
|
||||
paddle.pir.set_insertion_point_to_block_end(block)
|
||||
|
||||
|
||||
def _append_scale_op_after_comm(block, optimizer_ops, k_steps):
|
||||
for optimizer_op in optimizer_ops:
|
||||
target_value = None
|
||||
if "adam" in optimizer_op.name(): # adam and adamw are included
|
||||
target_value = optimizer_op.operand_source(1)
|
||||
elif "sgd" in optimizer_op.name():
|
||||
target_value = optimizer_op.operand_source(2)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"We yet support adamw, adam and sgd, but got {optimizer_op.name()}"
|
||||
)
|
||||
assert target_value is not None, (
|
||||
"target_value is not expected to be None"
|
||||
)
|
||||
insertion_point = target_value.get_defining_op()
|
||||
if insertion_point is None:
|
||||
# target_value is a gradient_merge_var, which hasn't defining_op
|
||||
# so we find the prev op of optimizer_op, inserting a scale op behind.
|
||||
insertion_point = _get_prev_op(block, optimizer_op)
|
||||
paddle.pir.set_insertion_point_after(insertion_point)
|
||||
_insert_scale_op_after(target_value, optimizer_op, 1.0 / k_steps)
|
||||
paddle.pir.set_insertion_point_to_block_end(block)
|
||||
|
||||
|
||||
def _pir_append_scale_op(program, new_params_to_grads, k_steps):
|
||||
block = program.global_block()
|
||||
optimizer_ops = _find_trivial_optimizer_ops(block)
|
||||
if len(optimizer_ops) > 0:
|
||||
_append_scale_op_after_comm(block, optimizer_ops, k_steps)
|
||||
else:
|
||||
_append_scale_op_before_comm(block, new_params_to_grads, k_steps)
|
||||
|
||||
|
||||
def _pir_parse_program(
|
||||
main_program,
|
||||
startup_program,
|
||||
params_grads,
|
||||
k_steps,
|
||||
avg,
|
||||
gradient_sync_after_accumulate,
|
||||
):
|
||||
# step1: append gradient merge backward op to main_program
|
||||
new_params_to_grads = _pir_append_gradient_merge_backward_op(
|
||||
main_program, startup_program, params_grads
|
||||
)
|
||||
|
||||
# step2: move back reduce op to backward stage
|
||||
if not gradient_sync_after_accumulate:
|
||||
_pir_move_reduce_to_backward_stage(main_program, params_grads)
|
||||
|
||||
# _pir_remove_cast_for_master_grad(main_program, params_grads)
|
||||
|
||||
# step3: append scale op
|
||||
if avg:
|
||||
_pir_append_scale_op(main_program, new_params_to_grads, k_steps)
|
||||
|
||||
|
||||
@register_pass("auto_parallel_gradient_merge_pass")
|
||||
class GradientMergePass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("k_steps", -1)
|
||||
self.set_attr("avg", True)
|
||||
self._in_pir_mode = paddle.base.framework.get_flags(
|
||||
"FLAGS_enable_pir_api"
|
||||
)["FLAGS_enable_pir_api"]
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("k_steps") < 1:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _type(self):
|
||||
return PassType.COMM_OPT
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
k_steps = self.get_attr("k_steps", -1)
|
||||
avg = self.get_attr("avg", False)
|
||||
params_grads = self.get_attr("params_grads")
|
||||
gradient_sync_after_accumulate = self.get_attr(
|
||||
"gradient_sync_after_accumulate", False
|
||||
)
|
||||
|
||||
if self._in_pir_mode:
|
||||
with paddle.static.program_guard(main_program, startup_program):
|
||||
_pir_parse_program(
|
||||
main_program,
|
||||
startup_program,
|
||||
params_grads,
|
||||
k_steps,
|
||||
avg,
|
||||
gradient_sync_after_accumulate,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"auto_parallel_gradient_merge_pass() only support PIR now."
|
||||
)
|
||||
@@ -0,0 +1,286 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
is_backward_op,
|
||||
is_gradient_clip_op,
|
||||
is_optimize_op,
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
||||
set_var_dist_attr,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import (
|
||||
OP_ROLE_KEY,
|
||||
OpRole,
|
||||
)
|
||||
from paddle.framework import core
|
||||
from paddle.static import program_guard
|
||||
|
||||
from ..utils.log_utils import get_logger
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle.base import Variable
|
||||
|
||||
_supported_optimizer_type = [
|
||||
"adam",
|
||||
"adamax",
|
||||
"adamw",
|
||||
"decayed_adagrad",
|
||||
"momentum",
|
||||
"dgc_momentum",
|
||||
"lars_momentum",
|
||||
"merged_momentum",
|
||||
"lamb",
|
||||
"sgd",
|
||||
]
|
||||
|
||||
logger = get_logger(logging.INFO, "MasterGradPass")
|
||||
|
||||
|
||||
def _is_master_grad_cast_op(block, op):
|
||||
op_name = op.type
|
||||
if op_name != "cast":
|
||||
return False
|
||||
input_names = op.input_arg_names
|
||||
output_names = op.output_arg_names
|
||||
|
||||
assert len(input_names) == 1
|
||||
assert len(output_names) == 1
|
||||
|
||||
input_var_name = input_names[0]
|
||||
|
||||
return (
|
||||
"@master_grad_fp16" in input_var_name
|
||||
or "@master_grad_bf16" in input_var_name
|
||||
)
|
||||
|
||||
|
||||
def get_output_in_varlist(op, var_names) -> list[str]:
|
||||
grad_names = []
|
||||
for output_name in op.output_arg_names:
|
||||
if output_name in var_names:
|
||||
grad_names.append(output_name)
|
||||
return grad_names
|
||||
|
||||
|
||||
@register_pass("auto_parallel_master_grad_pass")
|
||||
class MasterGradPass(PassBase):
|
||||
"""
|
||||
Use the high precision gradient to replace the low precision gradient in optimizer to avoid inf/nan values of low precision.
|
||||
The high precision gradient 'master grad' will be used by communication operator, `update_loss_scaling`, `GradClip` and `optimizer`.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
self._completer = self.get_attr("completer")
|
||||
dist_context = self.get_attr("dist_context")
|
||||
params_grads = self.get_attr("params_grads")
|
||||
logger.debug(f"Origin main_program: {main_program}")
|
||||
self._add_master_grad(main_program, params_grads, dist_context)
|
||||
self._regenerate_optimizer(
|
||||
main_program, startup_program, params_grads, dist_context
|
||||
)
|
||||
logger.debug(f"After main program: {main_program}")
|
||||
|
||||
def _add_cast_op(self, cur_block, grad_names: list[str], dist_context):
|
||||
grad_first_ids = OrderedDict()
|
||||
for idx, op in enumerate(cur_block.ops):
|
||||
if is_optimize_op(op):
|
||||
break
|
||||
elif is_backward_op(op):
|
||||
var_names = get_output_in_varlist(op, grad_names)
|
||||
for var_name in var_names:
|
||||
if var_name not in grad_first_ids:
|
||||
grad_first_ids[var_name] = idx
|
||||
# Communication operators such as 'allreduce_sum' use input var as output.
|
||||
else:
|
||||
pass
|
||||
|
||||
# insert cast op
|
||||
for grad_name, idx in reversed(grad_first_ids.items()):
|
||||
grad_var = cur_block.var(grad_name)
|
||||
if (
|
||||
grad_var.dtype == paddle.float16
|
||||
or grad_var.dtype == paddle.bfloat16
|
||||
):
|
||||
is_fp16 = grad_var.dtype == paddle.float16
|
||||
producer_op = cur_block.ops[idx]
|
||||
producer_op_dist_attr = (
|
||||
dist_context.get_op_dist_attr_for_program(producer_op)
|
||||
)
|
||||
assert producer_op_dist_attr is not None, (
|
||||
f"The op: '{producer_op}' should be distributed"
|
||||
)
|
||||
ref_output_dist_attr = (
|
||||
producer_op_dist_attr.get_output_dist_attr(grad_name)
|
||||
)
|
||||
assert ref_output_dist_attr is not None, (
|
||||
f"The output: '{grad_name}' should be distributed"
|
||||
)
|
||||
ref_mesh = ref_output_dist_attr.process_mesh
|
||||
ref_dims_mapping = ref_output_dist_attr.dims_mapping
|
||||
ref_chunk_id = producer_op_dist_attr.chunk_id
|
||||
grad_half_precision_name = (
|
||||
grad_name + '@master_grad_fp16'
|
||||
if is_fp16
|
||||
else grad_name + '@master_grad_bf16'
|
||||
)
|
||||
grad_half_precision = cur_block.create_var(
|
||||
name=grad_half_precision_name,
|
||||
dtype=grad_var.dtype,
|
||||
shape=grad_var.shape,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
set_var_dist_attr(
|
||||
dist_context,
|
||||
grad_half_precision,
|
||||
ref_dims_mapping,
|
||||
ref_mesh,
|
||||
chunk_id=ref_chunk_id,
|
||||
)
|
||||
|
||||
producer_op_dist_attr = (
|
||||
dist_context.get_op_dist_attr_for_program(producer_op)
|
||||
)
|
||||
origin_out_dims_mapping = (
|
||||
producer_op_dist_attr.get_output_dims_mapping(grad_name)
|
||||
)
|
||||
producer_op._rename_output(grad_name, grad_half_precision.name)
|
||||
producer_op_dist_attr.set_output_dims_mapping(
|
||||
grad_half_precision.name, origin_out_dims_mapping
|
||||
)
|
||||
grad_var.desc.set_dtype(core.VarDesc.VarType.FP32)
|
||||
|
||||
cast_op = cur_block._insert_op_without_sync(
|
||||
idx + 1,
|
||||
type="cast",
|
||||
inputs={"X": grad_half_precision},
|
||||
outputs={"Out": grad_var},
|
||||
attrs={
|
||||
"in_dtype": grad_half_precision.dtype,
|
||||
"out_dtype": grad_var.dtype,
|
||||
},
|
||||
)
|
||||
cast_op._set_attr(OP_ROLE_KEY, OpRole.Backward)
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
cast_op,
|
||||
ref_mesh,
|
||||
ref_dims_mapping,
|
||||
dist_context,
|
||||
chunk_id=ref_chunk_id,
|
||||
)
|
||||
cur_block._sync_with_cpp()
|
||||
|
||||
def _regenerate_optimizer(
|
||||
self,
|
||||
main_program,
|
||||
startup_program,
|
||||
params_grads: list[tuple[Variable, Variable]],
|
||||
dist_context,
|
||||
):
|
||||
grad_names = [g.name for _, g in params_grads]
|
||||
# 1. delete the origin optimizer op
|
||||
# 1.1 delete the var and op associated with the optimizer op in main_program
|
||||
main_ops = main_program.global_block().ops
|
||||
main_ops_len = len(main_ops)
|
||||
first_optimize_idx = main_ops_len
|
||||
for idx, op in enumerate(main_ops):
|
||||
# We don't delete the operators for check_nan_inf
|
||||
if is_optimize_op(op) and is_gradient_clip_op(op):
|
||||
first_optimize_idx = idx
|
||||
break
|
||||
assert first_optimize_idx < main_ops_len, (
|
||||
"The first optimizer op is not found!"
|
||||
)
|
||||
deleted_temp_var_names = []
|
||||
deleted_persist_var_names = []
|
||||
reserved_var_names = []
|
||||
for idx in range(main_ops_len - 1, first_optimize_idx - 1, -1):
|
||||
op = main_ops[idx]
|
||||
inout_arg_names = op.input_arg_names + op.output_arg_names
|
||||
if op.type in _supported_optimizer_type:
|
||||
param_names = op.input("Param")
|
||||
skip_update_names = op.input("SkipUpdate")
|
||||
for reserved_name in param_names + skip_update_names:
|
||||
if reserved_name not in reserved_var_names:
|
||||
reserved_var_names.append(reserved_name)
|
||||
for input_name in inout_arg_names:
|
||||
if input_name in grad_names:
|
||||
continue
|
||||
var = main_program.global_block().var(input_name)
|
||||
if (
|
||||
var.persistable
|
||||
and input_name not in deleted_persist_var_names
|
||||
):
|
||||
deleted_persist_var_names.append(input_name)
|
||||
elif (
|
||||
not var.persistable
|
||||
and input_name not in deleted_temp_var_names
|
||||
):
|
||||
deleted_temp_var_names.append(input_name)
|
||||
main_program.global_block()._remove_op(idx)
|
||||
|
||||
for var_name in deleted_temp_var_names + deleted_persist_var_names:
|
||||
if var_name not in reserved_var_names:
|
||||
main_program.global_block()._remove_var(var_name)
|
||||
main_program.global_block()._sync_with_cpp()
|
||||
|
||||
# 1.2 delete the var and op in startup_program
|
||||
for reserved_name in reserved_var_names:
|
||||
if reserved_name in deleted_persist_var_names:
|
||||
deleted_persist_var_names.remove(reserved_name)
|
||||
startup_global_block = startup_program.global_block()
|
||||
for var_name in deleted_persist_var_names:
|
||||
if startup_global_block.has_var(var_name):
|
||||
startup_global_block._remove_var(var_name)
|
||||
for idx, op in reversed(list(enumerate(startup_global_block.ops))):
|
||||
inout_arg_names = op.input_arg_names + op.output_arg_names
|
||||
for var_name in inout_arg_names:
|
||||
if var_name in deleted_persist_var_names:
|
||||
startup_program.global_block()._remove_op(idx)
|
||||
break
|
||||
|
||||
# 2. re-generate new optimizer op
|
||||
serial_optimizer = copy.deepcopy(dist_context._serial_optimizer)
|
||||
serial_optimizer._learning_rate = (
|
||||
dist_context._serial_optimizer._learning_rate
|
||||
)
|
||||
serial_optimizer._sorted = False
|
||||
with (
|
||||
program_guard(main_program, startup_program),
|
||||
main_program.switch_name_generator_guard("opt_"),
|
||||
):
|
||||
_ = serial_optimizer.apply_gradients(params_grads)
|
||||
self._completer.complete_update_annotation(main_program)
|
||||
|
||||
def _add_master_grad(self, main_program, params_grads, dist_context):
|
||||
grad_names = [g.name for _, g in params_grads]
|
||||
for sub_block in main_program.blocks:
|
||||
self._add_cast_op(sub_block, grad_names, dist_context)
|
||||
@@ -0,0 +1,463 @@
|
||||
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.framework import IrGraph, core
|
||||
from paddle.static.quantization import (
|
||||
AddQuantDequantForInferencePass,
|
||||
AddQuantDequantPassV2,
|
||||
OutScaleForTrainingPass,
|
||||
QuantizationTransformPassV2,
|
||||
quant_config,
|
||||
)
|
||||
|
||||
from ..auto_parallel.static.converter import Converter
|
||||
from ..auto_parallel.static.dist_attribute import (
|
||||
OperatorDistAttr,
|
||||
TensorDistAttr,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
TRANSFORM_PASS_OP_TYPES = list(
|
||||
quant_config.SUPPORT_WEIGHT_QUANTIZATION_OP_DICT.keys()
|
||||
)
|
||||
QUANT_DEQUANT_PASS_OP_TYPES = list(
|
||||
quant_config.SUPPORT_ACT_QUANTIZATION_OP_DICT.keys()
|
||||
)
|
||||
|
||||
|
||||
def _node_id(node):
|
||||
return (node.node.graph_id(), node.node.id())
|
||||
|
||||
|
||||
@register_pass("auto_parallel_quantization")
|
||||
class QuantizationPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("dist_context", None)
|
||||
self.set_attr("params_grads", None)
|
||||
self.set_attr("mode", "train")
|
||||
self.set_attr("loss", None)
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
if self.get_attr("params_grads") is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
dist_context = self.get_attr("dist_context")
|
||||
params_grads = self.get_attr("params_grads")
|
||||
mode = self.get_attr("mode")
|
||||
loss = self.get_attr("loss")
|
||||
|
||||
# TODO: scope and place will be removed,
|
||||
# cause params should be initialized by engine module.
|
||||
scope = paddle.static.global_scope()
|
||||
place = paddle.framework.CUDAPlace(
|
||||
paddle.distributed.ParallelEnv().dev_id
|
||||
)
|
||||
|
||||
# 0. record the relation among blocks
|
||||
parent_idx_dict = {}
|
||||
for block in main_program.blocks:
|
||||
parent_idx_dict[block.idx] = block.parent_idx
|
||||
|
||||
is_test = True if mode != "train" else False
|
||||
# 1. Program convert to Graph, and this pass is only for train mode
|
||||
main_graph = IrGraph(
|
||||
core.Graph(main_program.desc), for_test=mode != "train"
|
||||
)
|
||||
|
||||
# 2. Prepare inputs
|
||||
transform_pass_ops = []
|
||||
quant_dequant_ops = []
|
||||
quantize_op_types = [
|
||||
'conv2d',
|
||||
'depthwise_conv2d',
|
||||
'mul',
|
||||
'matmul',
|
||||
'matmul_v2',
|
||||
]
|
||||
for op_type in quantize_op_types:
|
||||
if op_type in TRANSFORM_PASS_OP_TYPES:
|
||||
transform_pass_ops.append(op_type)
|
||||
elif op_type in QUANT_DEQUANT_PASS_OP_TYPES:
|
||||
quant_dequant_ops.append(op_type)
|
||||
|
||||
weight_quantize_type = (
|
||||
"channel_wise_abs_max"
|
||||
if self.get_attr('channel_wise_abs_max')
|
||||
else "abs_max"
|
||||
)
|
||||
|
||||
# 3. Add quant op for ops which have parameters
|
||||
if len(transform_pass_ops) > 0:
|
||||
transform_pass = QuantizationTransformPassV2(
|
||||
scope=scope,
|
||||
place=place,
|
||||
weight_bits=self.get_attr('weight_bits'),
|
||||
activation_bits=self.get_attr('activation_bits'),
|
||||
skip_pattern=self.get_attr('not_quant_pattern'),
|
||||
activation_quantize_type="moving_average_abs_max",
|
||||
quantizable_op_type=transform_pass_ops,
|
||||
weight_quantize_type=weight_quantize_type,
|
||||
weight_quantize_func=None,
|
||||
act_quantize_func=None,
|
||||
weight_preprocess_func=None,
|
||||
act_preprocess_func=None,
|
||||
optimizer_func=None,
|
||||
executor=None,
|
||||
is_test=is_test,
|
||||
)
|
||||
for sub_graph in main_graph.all_sub_graphs():
|
||||
transform_pass.apply(sub_graph)
|
||||
|
||||
# 4. Add quant op for ops which don't have parameter
|
||||
if len(quant_dequant_ops) > 0:
|
||||
quant_dequant_pass = AddQuantDequantPassV2(
|
||||
scope=scope,
|
||||
place=place,
|
||||
quant_bits=self.get_attr('activation_bits'),
|
||||
skip_pattern=self.get_attr('not_quant_pattern'),
|
||||
quantizable_op_type=quant_dequant_ops,
|
||||
is_test=is_test,
|
||||
)
|
||||
for sub_graph in main_graph.all_sub_graphs():
|
||||
quant_dequant_pass.apply(sub_graph)
|
||||
|
||||
# 5. Gather quantitative information for the output
|
||||
out_scale_training_pass = OutScaleForTrainingPass(
|
||||
scope=scope, place=place, is_test=is_test
|
||||
)
|
||||
for sub_graph in main_graph.all_sub_graphs():
|
||||
out_scale_training_pass.apply(sub_graph)
|
||||
|
||||
# 6. When export quant model, traverse to find the output of each op, and insert the quant/dequant op after it.
|
||||
if mode != "train" and self.get_attr('onnx_format'):
|
||||
try:
|
||||
out_scale_infer_pass = AddQuantDequantForInferencePass(
|
||||
scope=scope,
|
||||
place=place,
|
||||
quant_bits=self.get_attr('activation_bits'),
|
||||
)
|
||||
# for sub_graph in main_graph.all_sub_graphs():
|
||||
# out_scale_infer_pass.apply(sub_graph)
|
||||
except:
|
||||
logging.warning(
|
||||
"Unable to convert quant model with onnx_format=True, please update PaddlePaddle >= 2.4.0"
|
||||
)
|
||||
|
||||
# 7. Convert Graph back to Program
|
||||
quant_program = main_graph.to_program()
|
||||
quant_program = self.move_persist_var_to_global_block(quant_program)
|
||||
|
||||
# 8.1 get new prams_grads from quant_program
|
||||
new_params_grads = []
|
||||
for param, grad in params_grads:
|
||||
if param.name not in quant_program.global_block().vars:
|
||||
continue
|
||||
|
||||
new_param = quant_program.global_block().vars[param.name]
|
||||
new_grad = quant_program.global_block().vars[grad.name]
|
||||
new_params_grads.append((new_param, new_grad))
|
||||
|
||||
# 8.2 get new loss var
|
||||
new_loss = None
|
||||
if loss:
|
||||
new_loss = quant_program.global_block().vars[loss.name]
|
||||
|
||||
# 8.3 recover the relation among blocks
|
||||
for block in quant_program.blocks:
|
||||
block.desc._set_forward_block_idx(parent_idx_dict[block.idx])
|
||||
|
||||
# 9. complete distributed attribution
|
||||
self.set_dist_attr_for_qat_program(
|
||||
quant_program, main_program, dist_context
|
||||
)
|
||||
|
||||
# 10. reset scale var value with dist_attr
|
||||
self.reset_scope_var(quant_program, dist_context, scope, place)
|
||||
|
||||
context.set_attr("main_program", quant_program)
|
||||
context.set_attr("startup_program", startup_program)
|
||||
context.set_attr("params_grads", new_params_grads)
|
||||
context.set_attr("loss", new_loss)
|
||||
|
||||
def move_persist_var_to_global_block(self, program):
|
||||
global_block = program.global_block()
|
||||
for _op in global_block.ops:
|
||||
if _op.type == "while":
|
||||
_block_id = _op.attr("sub_block").id
|
||||
_block = program.block(_block_id)
|
||||
persistables = []
|
||||
for _name, _var in _block.vars.items():
|
||||
if _var.persistable:
|
||||
global_block._clone_variable(_var)
|
||||
persistables.append(_name)
|
||||
for _name in persistables:
|
||||
_block._remove_var(_name)
|
||||
persistables.extend(_op.input('X'))
|
||||
_op.desc.set_input("X", persistables)
|
||||
return program
|
||||
|
||||
def reset_scope_var(self, quant_program, dist_context, scope, place):
|
||||
# The var_value, created by quantization_passes, should has same shape with the value after parallel.
|
||||
for var in quant_program.list_vars():
|
||||
scope_var = scope.find_var(var.name)
|
||||
if not (scope_var and scope_var.get_tensor()._is_initialized()):
|
||||
continue
|
||||
tensor = scope_var.get_tensor()
|
||||
if var.shape == tensor.shape:
|
||||
continue
|
||||
|
||||
var_dist_attr = dist_context.get_tensor_dist_attr_for_program(var)
|
||||
dist_attr = {
|
||||
"dims_mapping": var_dist_attr.dims_mapping,
|
||||
"process_shape": var_dist_attr.process_mesh.shape,
|
||||
"process_group": var_dist_attr.process_mesh.process_ids,
|
||||
}
|
||||
|
||||
# slice tensor_value with dist_attr
|
||||
sliced_tensor = Converter.slice_with_dist_attr(
|
||||
np.array(tensor), dist_attr
|
||||
)
|
||||
tensor._clear()
|
||||
tensor.set(sliced_tensor, place)
|
||||
|
||||
def set_dist_attr_for_qat_program(
|
||||
self, quant_program, main_program, dist_context
|
||||
):
|
||||
# NOTE: hack implement, upgrading soon
|
||||
for ib, block in enumerate(quant_program.blocks):
|
||||
# recover origin ops' dist_attr and set quant ops' dist_attr
|
||||
qat_offset = 0
|
||||
for ip, quant_op in enumerate(block.ops):
|
||||
quant_op_dist_attr = OperatorDistAttr()
|
||||
|
||||
if (
|
||||
"quantize" in quant_op.type
|
||||
or quant_op.type == "moving_average_abs_max_scale"
|
||||
):
|
||||
# set all quantization ops' dist_attr by quantified op
|
||||
input_name = quant_op.desc.input('X')[0]
|
||||
if "quantize" in input_name:
|
||||
input_name = input_name[
|
||||
: input_name.index(".quantized")
|
||||
]
|
||||
|
||||
if (
|
||||
quant_op.type == "moving_average_abs_max_scale"
|
||||
or ip - qat_offset >= len(main_program.blocks[ib].ops)
|
||||
):
|
||||
consume_op = (
|
||||
main_program.blocks[ib]
|
||||
._var_recursive(input_name)
|
||||
.op
|
||||
)
|
||||
else:
|
||||
consume_op = main_program.blocks[ib].ops[
|
||||
ip - qat_offset
|
||||
]
|
||||
consume_op_dist_attr = dist_context.get_dist_op_for_program(
|
||||
consume_op
|
||||
).dist_attr
|
||||
ref_process_mesh = consume_op_dist_attr.process_mesh
|
||||
|
||||
if input_name in consume_op_dist_attr.outputs_dist_attrs:
|
||||
consume_input_dist_attr = (
|
||||
consume_op_dist_attr.outputs_dist_attrs[input_name]
|
||||
)
|
||||
else:
|
||||
consume_input_dist_attr = (
|
||||
consume_op_dist_attr.inputs_dist_attrs[input_name]
|
||||
)
|
||||
|
||||
quant_op_dist_attr.impl_idx = 0
|
||||
quant_op_dist_attr.impl_type = "default"
|
||||
quant_op_dist_attr.process_mesh = ref_process_mesh
|
||||
quant_op_dist_attr.set_input_dist_attr(
|
||||
quant_op.desc.input('X')[0], consume_input_dist_attr
|
||||
)
|
||||
|
||||
for slot_name in quant_op.desc.input_names():
|
||||
in_name = quant_op.desc.input(slot_name)[0]
|
||||
input_var = block._var_recursive(in_name)
|
||||
ref_dims_mapping = [-1 for i in input_var.shape]
|
||||
if slot_name == "X":
|
||||
continue
|
||||
elif slot_name in ['Scale', 'ZeroPoint']:
|
||||
if (
|
||||
quant_op.has_attr('quant_axis')
|
||||
and quant_op.attr('quant_axis') != -1
|
||||
):
|
||||
x_name = quant_op.desc.input('X')[0]
|
||||
x_var = block._var_recursive(x_name)
|
||||
x_dist_attr = (
|
||||
quant_op_dist_attr.get_input_dist_attr(
|
||||
x_name
|
||||
)
|
||||
)
|
||||
quant_axis = quant_op.attr('quant_axis')
|
||||
ref_dims_mapping = [
|
||||
x_dist_attr.dims_mapping[quant_axis]
|
||||
]
|
||||
|
||||
tensor_dist_attr = TensorDistAttr()
|
||||
tensor_dist_attr.process_mesh = ref_process_mesh
|
||||
tensor_dist_attr.dims_mapping = ref_dims_mapping
|
||||
dist_context.set_tensor_dist_attr_for_program(
|
||||
input_var, tensor_dist_attr
|
||||
)
|
||||
quant_op_dist_attr.set_input_dist_attr(
|
||||
in_name, tensor_dist_attr
|
||||
)
|
||||
|
||||
for slot_name in quant_op.desc.output_names():
|
||||
output_name = quant_op.desc.output(slot_name)[0]
|
||||
output_var = block._var_recursive(output_name)
|
||||
ref_dims_mapping = [-1 for i in output_var.shape]
|
||||
if slot_name == "Y":
|
||||
dist_context.set_tensor_dist_attr_for_program(
|
||||
output_var, consume_input_dist_attr
|
||||
)
|
||||
quant_op_dist_attr.set_output_dist_attr(
|
||||
output_name, consume_input_dist_attr
|
||||
)
|
||||
continue
|
||||
elif slot_name == "OutScale":
|
||||
if (
|
||||
quant_op.has_attr('quant_axis')
|
||||
and quant_op.attr('quant_axis') != -1
|
||||
):
|
||||
x_name = quant_op.desc.input('X')[0]
|
||||
x_var = block._var_recursive(x_name)
|
||||
x_dist_attr = (
|
||||
quant_op_dist_attr.get_input_dist_attr(
|
||||
x_name
|
||||
)
|
||||
)
|
||||
quant_axis = quant_op.attr('quant_axis')
|
||||
ref_dims_mapping = [
|
||||
x_dist_attr.dims_mapping[quant_axis]
|
||||
]
|
||||
|
||||
tensor_dist_attr = TensorDistAttr()
|
||||
tensor_dist_attr.process_mesh = ref_process_mesh
|
||||
tensor_dist_attr.dims_mapping = ref_dims_mapping
|
||||
dist_context.set_tensor_dist_attr_for_program(
|
||||
output_var, tensor_dist_attr
|
||||
)
|
||||
quant_op_dist_attr.set_output_dist_attr(
|
||||
output_name, tensor_dist_attr
|
||||
)
|
||||
|
||||
quant_op._set_attr("op_device", "")
|
||||
qat_offset += 1
|
||||
|
||||
else:
|
||||
# recover origin ops' dist_attr
|
||||
origin_op = main_program.blocks[ib].ops[ip - qat_offset]
|
||||
quant_op.desc.set_original_id(origin_op.desc.original_id())
|
||||
dist_origin_op = dist_context.get_dist_op_for_program(
|
||||
origin_op
|
||||
)
|
||||
assert dist_origin_op is not None, (
|
||||
"origin op must have dist attr."
|
||||
)
|
||||
|
||||
origin_op_dist_attr = dist_origin_op.dist_attr
|
||||
quant_op_dist_attr.impl_idx = origin_op_dist_attr.impl_idx
|
||||
quant_op_dist_attr.impl_type = origin_op_dist_attr.impl_type
|
||||
quant_op_dist_attr.process_mesh = (
|
||||
origin_op_dist_attr.process_mesh
|
||||
)
|
||||
|
||||
scale_offset = 0
|
||||
for idx, input_name in enumerate(quant_op.input_arg_names):
|
||||
if (
|
||||
origin_op.type == "while"
|
||||
and input_name not in origin_op.input_arg_names
|
||||
):
|
||||
assert (
|
||||
"@scale" in input_name
|
||||
or "@zero_point" in input_name
|
||||
)
|
||||
scale_offset += 1
|
||||
continue
|
||||
|
||||
idx -= scale_offset
|
||||
origin_input_name = origin_op.input_arg_names[idx]
|
||||
origin_input_dist_attr = (
|
||||
origin_op_dist_attr.inputs_dist_attrs[
|
||||
origin_input_name
|
||||
]
|
||||
)
|
||||
quant_op_dist_attr.set_input_dist_attr(
|
||||
input_name, origin_input_dist_attr
|
||||
)
|
||||
|
||||
for idx, output_name in enumerate(
|
||||
quant_op.output_arg_names
|
||||
):
|
||||
origin_output_name = origin_op.output_arg_names[idx]
|
||||
origin_output_dist_attr = (
|
||||
origin_op_dist_attr.outputs_dist_attrs[
|
||||
origin_output_name
|
||||
]
|
||||
)
|
||||
quant_op_dist_attr.set_output_dist_attr(
|
||||
output_name, origin_output_dist_attr
|
||||
)
|
||||
|
||||
if not main_program.blocks[ib]._find_var_recursive(
|
||||
output_name
|
||||
):
|
||||
origin_output_var = main_program.blocks[
|
||||
ib
|
||||
]._var_recursive(origin_output_name)
|
||||
origin_out_tensor_dist_attr = (
|
||||
dist_context.get_dist_tensor_for_program(
|
||||
origin_output_var
|
||||
).dist_attr
|
||||
)
|
||||
quant_output_var = block._var_recursive(output_name)
|
||||
dist_context.set_tensor_dist_attr_for_program(
|
||||
quant_output_var, origin_out_tensor_dist_attr
|
||||
)
|
||||
|
||||
dist_context.set_op_dist_attr_for_program(
|
||||
quant_op, quant_op_dist_attr
|
||||
)
|
||||
|
||||
# recover vars' dist_attr
|
||||
for name, dst_var in block.vars.items():
|
||||
if name in main_program.blocks[ib].vars:
|
||||
src_var = main_program.blocks[ib].vars[name]
|
||||
dist_tensor = dist_context.get_dist_tensor_for_program(
|
||||
src_var
|
||||
)
|
||||
if not dist_tensor:
|
||||
continue
|
||||
dist_context.set_tensor_dist_attr_for_program(
|
||||
dst_var, dist_tensor.dist_attr
|
||||
)
|
||||
@@ -0,0 +1,628 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
import re
|
||||
|
||||
import paddle
|
||||
from paddle.base.backward import (
|
||||
ProgramStats,
|
||||
_append_grad_suffix_,
|
||||
_find_op_path_,
|
||||
_get_no_grad_set_name,
|
||||
_rename_arg_,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
|
||||
from paddle.framework import core
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ..auto_parallel.static.dist_attribute import OperatorDistAttr
|
||||
from ..auto_parallel.static.utils import (
|
||||
get_loss_op,
|
||||
insert_dependencies_for_two_ops,
|
||||
is_backward_op,
|
||||
is_recompute_exclude_op,
|
||||
is_recompute_op,
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
||||
set_dist_op_desc_original_id,
|
||||
set_var_dist_attr,
|
||||
)
|
||||
from ..utils.log_utils import get_logger
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
class RecomputeState(ProgramStats):
|
||||
def __init__(self, block, ops):
|
||||
super().__init__(block=block, ops=ops)
|
||||
self.seg_op_deps = {}
|
||||
self._checkpoints = []
|
||||
self._reserved_vars = []
|
||||
|
||||
@property
|
||||
def checkpoints(self):
|
||||
return self._checkpoints
|
||||
|
||||
@property
|
||||
def reserved_vars(self):
|
||||
return self._reserved_vars
|
||||
|
||||
def is_recompute(self):
|
||||
return any(is_recompute_op(op) for op in self.ops)
|
||||
|
||||
def build_states(self):
|
||||
for i, op in enumerate(self.ops):
|
||||
if is_backward_op(op):
|
||||
break
|
||||
|
||||
for name in op.input_arg_names:
|
||||
if name in self.var_op_deps:
|
||||
self.var_op_deps[name]["var_as_input_ops"].extend([i])
|
||||
else:
|
||||
self.var_op_deps[name] = {}
|
||||
self.var_op_deps[name]["var_as_input_ops"] = [i]
|
||||
self.var_op_deps[name]["var_as_output_ops"] = []
|
||||
|
||||
for name in op.output_arg_names:
|
||||
if name in self.var_op_deps:
|
||||
self.var_op_deps[name]["var_as_output_ops"].extend([i])
|
||||
else:
|
||||
self.var_op_deps[name] = {}
|
||||
self.var_op_deps[name]["var_as_input_ops"] = []
|
||||
self.var_op_deps[name]["var_as_output_ops"] = [i]
|
||||
|
||||
if not is_recompute_op(op):
|
||||
self._checkpoints.extend(op.output_arg_names)
|
||||
if not is_recompute_exclude_op(op):
|
||||
continue
|
||||
|
||||
seg_name = op.attr('op_namescope')
|
||||
res = re.search("/auto_parallel/rc_[0-9]*", seg_name)
|
||||
seg_name = res.group(0)
|
||||
if seg_name not in self.seg_op_deps:
|
||||
self.seg_op_deps[seg_name] = [i]
|
||||
else:
|
||||
assert self.seg_op_deps[seg_name][-1] + 1 == i, (
|
||||
"The recompute segment's ops should be continuous"
|
||||
)
|
||||
self.seg_op_deps[seg_name].extend([i])
|
||||
|
||||
def get_recompute_segments(self, no_recompute_segments=[]):
|
||||
segments = []
|
||||
for segment_idx in self.seg_op_deps.values():
|
||||
if len(segment_idx) == 1:
|
||||
continue
|
||||
segments.append([segment_idx[0], segment_idx[-1] + 1])
|
||||
self._checkpoints.extend(self.ops[segment_idx[-1]].output_arg_names)
|
||||
|
||||
for i in sorted(no_recompute_segments, reverse=True):
|
||||
assert i < len(segments), (
|
||||
f"the no_recompute_segments idx [{i}] should be lower the number of segment [{len(segments)}]"
|
||||
)
|
||||
segments.pop(i)
|
||||
|
||||
return segments
|
||||
|
||||
def modify_forward_desc_for_recompute(self, dist_context):
|
||||
"""
|
||||
If program's forward part has 'dropout' op, this function will insert
|
||||
a seed op before it to guarantee that two dropout op have the same outputs.
|
||||
"""
|
||||
op_types = [op.type for op in self.ops]
|
||||
if "dropout" not in op_types and "fused_dropout_add" not in op_types:
|
||||
return
|
||||
|
||||
op_idx = 0
|
||||
while op_idx < len(self.ops):
|
||||
cur_op = self.ops[op_idx]
|
||||
if "grad" in cur_op.type:
|
||||
break
|
||||
if cur_op.type == "seed":
|
||||
self._reserved_vars.extend(cur_op.output_arg_names)
|
||||
op_idx += 1
|
||||
continue
|
||||
if cur_op.type not in ["dropout", "fused_dropout_add"]:
|
||||
op_idx += 1
|
||||
continue
|
||||
seed_tensor_name = (
|
||||
"seed_tensor" if cur_op.type == "fused_dropout_add" else "Seed"
|
||||
)
|
||||
if cur_op.input(seed_tensor_name) is not None and len(
|
||||
cur_op.input(seed_tensor_name)
|
||||
):
|
||||
op_idx += 1
|
||||
continue
|
||||
|
||||
cur_op_dist_attr = dist_context.get_op_dist_attr_for_program(cur_op)
|
||||
# insert seed op to guarantee that two dropout op have the same outputs
|
||||
# NOTE Hack for adopt recompute for random control, for more info see dist_dropout.py
|
||||
# new seed added by recompute should have a prefix to distinguish with seed added by user or other module.
|
||||
op_unique_name = unique_name.generate("rc_seed")
|
||||
var_unique_name = unique_name.generate_with_ignorable_key(
|
||||
".".join([op_unique_name, 'tmp'])
|
||||
)
|
||||
self._reserved_vars.append(var_unique_name)
|
||||
seed_var = self.block.create_var(
|
||||
name=var_unique_name,
|
||||
dtype='int32',
|
||||
type=core.VarDesc.VarType.DENSE_TENSOR,
|
||||
persistable=False,
|
||||
stop_gradient=False,
|
||||
)
|
||||
|
||||
# set new seed_var's dist_attr
|
||||
ref_dims_mapping = [-1]
|
||||
ref_process_mesh = cur_op_dist_attr.process_mesh
|
||||
seed_var_dist_attr = set_var_dist_attr(
|
||||
dist_context,
|
||||
seed_var,
|
||||
ref_dims_mapping,
|
||||
ref_process_mesh,
|
||||
chunk_id=cur_op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
seed = (
|
||||
0
|
||||
if cur_op.attr("fix_seed") is False
|
||||
else int(cur_op.attr("seed"))
|
||||
)
|
||||
# TODO add dependency for seed op to ensure it be issued just before recompute.
|
||||
seed_op = self.block._insert_op_without_sync(
|
||||
index=cur_op.idx,
|
||||
type="seed",
|
||||
inputs={},
|
||||
outputs={"Out": seed_var},
|
||||
attrs={"seed": seed, "force_cpu": True},
|
||||
)
|
||||
seed_op._set_attr('op_namescope', cur_op.attr('op_namescope'))
|
||||
# set new seed op's dist_attr
|
||||
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
||||
seed_op,
|
||||
ref_process_mesh,
|
||||
ref_dims_mapping,
|
||||
dist_context,
|
||||
chunk_id=cur_op_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# modify dropout op's desc
|
||||
self.ops.insert(op_idx, seed_op)
|
||||
cur_op.desc.set_input(seed_tensor_name, [var_unique_name])
|
||||
cur_op.desc._set_attr("fix_seed", False)
|
||||
cur_op.desc._set_attr("seed", 0)
|
||||
cur_op_dist_attr.set_input_dist_attr(
|
||||
seed_var.name, seed_var_dist_attr
|
||||
)
|
||||
op_idx += 2
|
||||
|
||||
self.block._sync_with_cpp()
|
||||
|
||||
|
||||
def _find_op_index(block, cur_op):
|
||||
for idx in range(block.desc.op_size()):
|
||||
if cur_op.desc == block.desc.op(idx):
|
||||
return idx
|
||||
return -1
|
||||
|
||||
|
||||
def _get_stop_gradients(program, no_grad_set=None):
|
||||
"""get no grad var"""
|
||||
if no_grad_set is None:
|
||||
no_grad_set = set()
|
||||
else:
|
||||
no_grad_set = _get_no_grad_set_name(no_grad_set)
|
||||
|
||||
no_grad_set_name = set()
|
||||
for var in program.list_vars():
|
||||
if "@GRAD" in var.name:
|
||||
break
|
||||
if var.stop_gradient:
|
||||
no_grad_set_name.add(_append_grad_suffix_(var.name))
|
||||
no_grad_set_name.update(list(map(_append_grad_suffix_, no_grad_set)))
|
||||
return no_grad_set_name
|
||||
|
||||
|
||||
def _add_needed_descs_to_block(
|
||||
descs, block, main_block, vars_should_be_hold, dist_context
|
||||
):
|
||||
"""
|
||||
Get the recomputed ops which will insert the backward part
|
||||
"""
|
||||
if len(descs) == 0:
|
||||
return []
|
||||
|
||||
result_descs = []
|
||||
for desc in descs:
|
||||
# if isinstance(desc, framework.Operator):
|
||||
if isinstance(desc, paddle.static.Operator):
|
||||
desc = desc.desc
|
||||
if isinstance(desc, tuple):
|
||||
desc = desc[0]
|
||||
is_needed = False
|
||||
for name in desc.output_arg_names():
|
||||
if main_block.has_var(name) and main_block.var(name).persistable:
|
||||
continue
|
||||
if name not in vars_should_be_hold:
|
||||
is_needed = True
|
||||
if is_needed:
|
||||
new_op_desc = block.desc.append_op()
|
||||
new_op_desc.copy_from(desc)
|
||||
set_dist_op_desc_original_id(new_op_desc, desc, dist_context)
|
||||
new_op_desc._set_attr(OP_ROLE_KEY, OpRole.Backward)
|
||||
result_descs.append(new_op_desc)
|
||||
return result_descs
|
||||
|
||||
|
||||
def _find_op_path(main_program, loss, no_grad_set=None):
|
||||
no_grad_set_name = _get_stop_gradients(main_program, no_grad_set)
|
||||
op_path = _find_op_path_(
|
||||
main_program.global_block(), [loss], [], no_grad_set_name
|
||||
)
|
||||
return op_path
|
||||
|
||||
|
||||
@register_pass("auto_parallel_recompute")
|
||||
class RecomputePass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("loss", None)
|
||||
self.set_attr("dist_context", None)
|
||||
self.set_attr("no_grad_set", None)
|
||||
self.set_attr("no_recompute_segments", [])
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
if self.get_attr("loss") is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def get_ops_per_device(self, ops, all_ops_process_meshes, sr=0):
|
||||
"""
|
||||
Get ops and op_names of each process mesh excluding ops within the first "sr" chunks
|
||||
"""
|
||||
|
||||
def reset_recompute_op(op):
|
||||
if is_recompute_op(op) or is_recompute_exclude_op(op):
|
||||
op._set_attr("op_namescope", "")
|
||||
|
||||
all_process_meshes_count = len(all_ops_process_meshes)
|
||||
ops_of_stages = [[] for _ in range(all_process_meshes_count)]
|
||||
op_names_of_stages = [[] for _ in range(all_process_meshes_count)]
|
||||
pushed_ops_count = 0
|
||||
reset_ops_count = 0
|
||||
chunk_id = 0
|
||||
for op_id, op in enumerate(ops):
|
||||
if chunk_id // all_process_meshes_count < sr:
|
||||
reset_ops_count += 1
|
||||
reset_recompute_op(op)
|
||||
if (
|
||||
op_id < len(ops) - 1
|
||||
and op.dist_attr.process_mesh
|
||||
!= ops[op_id + 1].dist_attr.process_mesh
|
||||
):
|
||||
chunk_id += 1
|
||||
if chunk_id // all_process_meshes_count < sr:
|
||||
continue
|
||||
|
||||
for id, process_mesh in enumerate(all_ops_process_meshes):
|
||||
if op.dist_attr.process_mesh == process_mesh:
|
||||
pushed_ops_count += 1
|
||||
ops_of_stages[id].append(op)
|
||||
op_names_of_stages[id].append(op.type)
|
||||
assert len(ops) == reset_ops_count + pushed_ops_count, (
|
||||
f"The sum of pushed_ops_count and reset_ops_count must be the same as length of ops, but the sum is {reset_ops_count + pushed_ops_count} while length of ops is {len(ops)}"
|
||||
)
|
||||
return ops_of_stages, op_names_of_stages
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
loss = self.get_attr("loss")
|
||||
no_grad_set = self.get_attr("no_grad_set")
|
||||
no_recompute_segments = self.get_attr("no_recompute_segments")
|
||||
self._dist_context = self.get_attr("dist_context")
|
||||
self._sr = self.get_attr("sr", 0)
|
||||
self._refined_ops_patterns = self.get_attr("refined_ops_patterns", [])
|
||||
|
||||
# 0. get op_path which is related to loss
|
||||
main_block = main_program.global_block()
|
||||
op_path = _find_op_path(main_program, loss, no_grad_set)
|
||||
|
||||
# 1. mark exclude ops for refined-recompute according to ops-patterns(mainly linear and flash_attn)
|
||||
# 1.1 get all process_meshes in op_path
|
||||
all_ops_process_meshes = []
|
||||
for op in op_path:
|
||||
if op.dist_attr.process_mesh not in all_ops_process_meshes:
|
||||
all_ops_process_meshes.append(op.dist_attr.process_mesh)
|
||||
|
||||
# 1.2 get ops_devices and op_names_devices
|
||||
ops_devices, op_names_devices = self.get_ops_per_device(
|
||||
op_path, all_ops_process_meshes, self._sr
|
||||
)
|
||||
all_ops_len = len(op_path)
|
||||
all_exclude_ops_ids = [[] for _ in op_names_devices]
|
||||
# 1.3 find exclude ops for refined-recompute according to ops-patterns
|
||||
for refined_ops_pattern in self._refined_ops_patterns:
|
||||
num = refined_ops_pattern['num']
|
||||
num = (
|
||||
num if num >= 0 else all_ops_len
|
||||
) # 'num == -1' represents to all ops
|
||||
main_ops = refined_ops_pattern['main_ops']
|
||||
pre_ops = refined_ops_pattern['pre_ops']
|
||||
suf_ops = refined_ops_pattern['suf_ops']
|
||||
main_start_id = len(pre_ops)
|
||||
main_ops_len = len(main_ops)
|
||||
pattern_ops = pre_ops + main_ops + suf_ops
|
||||
pattern_ops_len = len(pattern_ops)
|
||||
|
||||
for id, op_names_device in enumerate(op_names_devices):
|
||||
pattern_count = 0
|
||||
ops_len_device = len(op_names_device)
|
||||
for i in range(ops_len_device - pattern_ops_len + 1):
|
||||
if (
|
||||
op_names_device[i : i + pattern_ops_len] == pattern_ops
|
||||
and pattern_count < num
|
||||
):
|
||||
pattern_count += 1
|
||||
all_exclude_ops_ids[id].extend(
|
||||
list(
|
||||
range(
|
||||
i + main_start_id,
|
||||
i + main_start_id + main_ops_len,
|
||||
)
|
||||
)
|
||||
)
|
||||
logger.info(
|
||||
f"The excluded ops in recompute segments are:\n{all_exclude_ops_ids}"
|
||||
)
|
||||
# 1.4 mark exclude ops in exclude_ops_ids
|
||||
for id, exclude_ops_ids in enumerate(all_exclude_ops_ids):
|
||||
for op_id in exclude_ops_ids:
|
||||
if is_recompute_op(ops_devices[id][op_id]):
|
||||
rc_mark_str = ops_devices[id][op_id].attr("op_namescope")
|
||||
ops_devices[id][op_id]._set_attr(
|
||||
"op_namescope", rc_mark_str + "_exclude_rc"
|
||||
)
|
||||
|
||||
# 2. build recompute state
|
||||
rc_state = RecomputeState(main_block, op_path)
|
||||
if not rc_state.is_recompute():
|
||||
return
|
||||
|
||||
# 3. get the segments to be recomputed
|
||||
rc_state.modify_forward_desc_for_recompute(self._dist_context)
|
||||
rc_state.build_states()
|
||||
segments = rc_state.get_recompute_segments(no_recompute_segments)
|
||||
if segments == []:
|
||||
return
|
||||
|
||||
for i, (idx1, idx2) in enumerate(segments):
|
||||
logger.debug(f"recompute segment[{i + 1}/{len(segments)}]")
|
||||
logger.debug(
|
||||
f"segment start op: [{rc_state.ops[idx1].type}]: [{rc_state.ops[idx1].input_arg_names}] [{rc_state.ops[idx1].output_arg_names}]"
|
||||
)
|
||||
logger.debug(
|
||||
f"segment end op: [{rc_state.ops[idx2 - 1].type}]: [{rc_state.ops[idx2 - 1].input_arg_names}] [{rc_state.ops[idx2 - 1].output_arg_names}]"
|
||||
)
|
||||
|
||||
# 4. get vars that should be hold in memory
|
||||
# list of var_names
|
||||
vars_should_be_hold = []
|
||||
for segment in segments:
|
||||
vars_should_be_hold.extend(
|
||||
rc_state.get_out_of_subgraph_vars(segment[0], segment[1])
|
||||
)
|
||||
cross_vars = set(vars_should_be_hold) - set(rc_state.checkpoints)
|
||||
logger.debug(
|
||||
f"found [{len(cross_vars)}] vars which cross recompute segment: [{cross_vars}],"
|
||||
"better checkpoints might be set to reduce those vars"
|
||||
)
|
||||
vars_should_be_hold.extend(rc_state.reserved_vars)
|
||||
vars_should_be_hold.extend(rc_state.get_input_nodes())
|
||||
vars_should_be_hold = list(
|
||||
set(vars_should_be_hold) | set(rc_state.checkpoints)
|
||||
)
|
||||
|
||||
# 5. get the fwd ops desc to be recomputed.
|
||||
var_name_dict = {} # varname --> varname.subprog_XXX
|
||||
ckpt_ops_dict = {} # ckpt_op_id --> segment_descs
|
||||
buffer_block = main_block.program._create_block()
|
||||
for i, segment in enumerate(segments[::-1]):
|
||||
fwd_ops = op_path[segment[0] : segment[1]]
|
||||
var_suffix = f".subprog_{i}"
|
||||
for op in fwd_ops:
|
||||
input_and_output_names = []
|
||||
input_and_output_names.extend(op.input_arg_names)
|
||||
input_and_output_names.extend(op.output_arg_names)
|
||||
|
||||
cur_op_dist_attr = (
|
||||
self._dist_context.get_op_dist_attr_for_program(op)
|
||||
)
|
||||
assert cur_op_dist_attr is not None
|
||||
|
||||
for name in input_and_output_names:
|
||||
if (
|
||||
main_block.var(name).persistable
|
||||
or name in vars_should_be_hold
|
||||
):
|
||||
continue
|
||||
if name not in var_name_dict:
|
||||
ref_process_mesh = cur_op_dist_attr.process_mesh
|
||||
if name in op.input_arg_names:
|
||||
ref_dims_mapping = (
|
||||
cur_op_dist_attr.get_input_dims_mapping(name)
|
||||
)
|
||||
else:
|
||||
ref_dims_mapping = (
|
||||
cur_op_dist_attr.get_output_dims_mapping(name)
|
||||
)
|
||||
|
||||
# record recomputed var's old_name and new_name (old_name.subprog_XXX)
|
||||
# create new var with new name
|
||||
var_name_dict[name] = name + var_suffix
|
||||
ref_var = main_block.var(name)
|
||||
rc_var = main_block.create_var(
|
||||
name=var_name_dict[name],
|
||||
shape=ref_var.shape,
|
||||
dtype=ref_var.dtype,
|
||||
type=ref_var.type,
|
||||
persistable=ref_var.persistable,
|
||||
stop_gradient=ref_var.stop_gradient,
|
||||
)
|
||||
# set new recomputed var's dist attr
|
||||
set_var_dist_attr(
|
||||
self._dist_context,
|
||||
rc_var,
|
||||
ref_dims_mapping,
|
||||
ref_process_mesh,
|
||||
chunk_id=cur_op_dist_attr.chunk_id,
|
||||
)
|
||||
# get recomputed segment's descs
|
||||
segment_descs = _add_needed_descs_to_block(
|
||||
fwd_ops,
|
||||
buffer_block,
|
||||
main_block,
|
||||
vars_should_be_hold,
|
||||
self._dist_context,
|
||||
)
|
||||
# rename recomputed ops' input and output var name
|
||||
for key in var_name_dict:
|
||||
_rename_arg_(segment_descs, key, var_name_dict[key])
|
||||
|
||||
# NOTE: one forward op could be correspond to multiple xxx_grad op.
|
||||
# When traversing all grad_ops in reverse, need to set a flag to indicate
|
||||
# whether the ckpt and its segment_descs can be used.
|
||||
ckpt_op = op_path[segment[1] - 1]
|
||||
ckpt_ops_dict[ckpt_op.desc.original_id()] = [True, segment_descs]
|
||||
|
||||
# 6. insert recomputed fwd ops into backward parse
|
||||
ops = main_block.ops
|
||||
loss_op = get_loss_op(main_block)
|
||||
loss_op_idx = _find_op_index(main_block, loss_op)
|
||||
dist_op_context = self._dist_context.dist_op_context
|
||||
assert loss_op_idx != -1
|
||||
# Traversing all grad_ops in reverse, and if the fwd op corresponding to reverse op is checkpoints,
|
||||
# segments ops should be inserted.
|
||||
for i in range(len(ops) - 1, loss_op_idx, -1):
|
||||
grad_op = ops[i]
|
||||
|
||||
input_and_output_names = []
|
||||
input_and_output_names.extend(grad_op.input_arg_names)
|
||||
input_and_output_names.extend(grad_op.output_arg_names)
|
||||
|
||||
for varname in var_name_dict:
|
||||
if varname not in input_and_output_names:
|
||||
continue
|
||||
self.reset_op_dist_attr(grad_op, var_name_dict)
|
||||
_rename_arg_([grad_op.desc], varname, var_name_dict[varname])
|
||||
|
||||
# insert recomputed ops
|
||||
original_id = grad_op.desc.original_id()
|
||||
if original_id in dist_op_context.grad_op_id_to_op_id:
|
||||
fwd_op_id = dist_op_context.grad_op_id_to_op_id[original_id]
|
||||
if fwd_op_id in ckpt_ops_dict and ckpt_ops_dict[fwd_op_id][0]:
|
||||
idx = grad_op.idx
|
||||
while idx - 1 >= 0 and ops[idx - 1].type == "sum":
|
||||
idx -= 1
|
||||
segment_descs = ckpt_ops_dict[fwd_op_id][1]
|
||||
rc_op = None
|
||||
for _, op_desc in reversed(list(enumerate(segment_descs))):
|
||||
rc_op = main_block._insert_op_without_sync(
|
||||
idx, type='nop'
|
||||
)
|
||||
rc_desc = rc_op.desc
|
||||
rc_desc.copy_from(op_desc)
|
||||
rc_desc.set_original_id(rc_desc.id())
|
||||
# set recomputed ops' dist attr
|
||||
fwd_op_dist_attr = self._dist_context.get_op_dist_attr_for_program_with_id(
|
||||
op_desc.original_id()
|
||||
)
|
||||
assert fwd_op_dist_attr is not None
|
||||
self.set_op_dist_attr(
|
||||
rc_op, fwd_op_dist_attr, var_name_dict
|
||||
)
|
||||
|
||||
ckpt_ops_dict[fwd_op_id][0] = False
|
||||
if rc_op:
|
||||
prior_op = main_block.ops[rc_op.idx - 1]
|
||||
posterior_op = rc_op
|
||||
prior_mesh = (
|
||||
self._dist_context.get_op_dist_attr_for_program(
|
||||
prior_op
|
||||
).process_mesh
|
||||
)
|
||||
posterior_mesh = (
|
||||
self._dist_context.get_op_dist_attr_for_program(
|
||||
posterior_op
|
||||
).process_mesh
|
||||
)
|
||||
# NOTE if two recompute segments across two pipeline stages
|
||||
# not need dependencies for it
|
||||
if prior_mesh == posterior_mesh:
|
||||
insert_dependencies_for_two_ops(
|
||||
main_block,
|
||||
idx,
|
||||
prior_op,
|
||||
posterior_op,
|
||||
self._dist_context,
|
||||
is_recompute=True,
|
||||
sync=False,
|
||||
op_namescope="recompute_segment_dep",
|
||||
)
|
||||
main_program._sync_with_cpp()
|
||||
|
||||
def reset_op_dist_attr(self, op, var_name_dict):
|
||||
op_dist_attr = self._dist_context.get_op_dist_attr_for_program(op)
|
||||
assert op_dist_attr is not None
|
||||
for input in op.input_arg_names:
|
||||
if input in var_name_dict.keys():
|
||||
in_dist_attr = op_dist_attr.get_input_dist_attr(input)
|
||||
op_dist_attr.set_input_dist_attr(
|
||||
var_name_dict[input], in_dist_attr
|
||||
)
|
||||
for output in op.output_arg_names:
|
||||
if output in var_name_dict.keys():
|
||||
out_dist_attr = op_dist_attr.get_output_dist_attr(output)
|
||||
op_dist_attr.set_output_dist_attr(
|
||||
var_name_dict[output], out_dist_attr
|
||||
)
|
||||
|
||||
def set_op_dist_attr(self, op, old_dist_attr, var_name_dict):
|
||||
new_dist_attr = OperatorDistAttr()
|
||||
new_dist_attr.is_recompute = True
|
||||
new_dist_attr.impl_idx = old_dist_attr.impl_idx
|
||||
new_dist_attr.impl_type = old_dist_attr.impl_type
|
||||
new_dist_attr.process_mesh = old_dist_attr.process_mesh
|
||||
new_dist_attr.chunk_id = old_dist_attr.chunk_id
|
||||
for input in old_dist_attr.inputs_dist_attrs.keys():
|
||||
if input in var_name_dict.keys():
|
||||
in_dist_attr = old_dist_attr.inputs_dist_attrs[input]
|
||||
new_dist_attr.set_input_dist_attr(
|
||||
var_name_dict[input], in_dist_attr
|
||||
)
|
||||
else:
|
||||
in_dist_attr = old_dist_attr.inputs_dist_attrs[input]
|
||||
new_dist_attr.set_input_dist_attr(input, in_dist_attr)
|
||||
for output in old_dist_attr.outputs_dist_attrs.keys():
|
||||
if output in var_name_dict.keys():
|
||||
out_dist_attr = old_dist_attr.outputs_dist_attrs[output]
|
||||
new_dist_attr.set_output_dist_attr(
|
||||
var_name_dict[output], out_dist_attr
|
||||
)
|
||||
else:
|
||||
out_dist_attr = old_dist_attr.outputs_dist_attrs[output]
|
||||
new_dist_attr.set_output_dist_attr(output, out_dist_attr)
|
||||
self._dist_context.set_op_dist_attr_for_program(op, new_dist_attr)
|
||||
@@ -0,0 +1,289 @@
|
||||
# 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
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
OpRole = core.op_proto_and_checker_maker.OpRole
|
||||
|
||||
from paddle.autograd import backward_utils
|
||||
|
||||
from ..auto_parallel.static.utils import (
|
||||
get_logger,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
@register_pass("auto_parallel_recompute_pir")
|
||||
class AutoParallelRecomputePIRPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.program_ops = []
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def get_fwd_bwd_ops(self):
|
||||
fwd_ops = []
|
||||
bwd_ops = []
|
||||
for op in self.program_ops:
|
||||
if op.op_role == int(OpRole.Forward):
|
||||
fwd_ops.append(op)
|
||||
elif op.op_role == int(OpRole.Backward):
|
||||
bwd_ops.append(op)
|
||||
assert len(fwd_ops) and len(bwd_ops)
|
||||
return fwd_ops, bwd_ops
|
||||
|
||||
def get_first_bwd_used_op(self, fwd_op, bwd_ops):
|
||||
# Find the first user op of the op result in backward op list.
|
||||
first_op = bwd_ops[-1]
|
||||
for res in fwd_op.results():
|
||||
for user_op in res.all_used_ops():
|
||||
if user_op in bwd_ops and first_op.id() >= user_op.id():
|
||||
first_op = user_op
|
||||
return first_op
|
||||
|
||||
def is_seed_used_by_dropout(self, seed_op):
|
||||
# Ensure that the random operator has the same output in backward recompute.
|
||||
if seed_op.name() != "seed":
|
||||
return False
|
||||
seed_value = seed_op.results()[0]
|
||||
dropout_ops = ["pd_op.dropout", "pd_op.fused_dropout_add"]
|
||||
return any(
|
||||
True
|
||||
for used_op in seed_value.all_used_ops()
|
||||
if used_op.name() in dropout_ops
|
||||
)
|
||||
|
||||
def remove_outgoing_op(self, segment):
|
||||
# An OP is considered an outgoing OP if all of results' user OPs are not in segment.
|
||||
# These OPs do not participate in the backward gradient computation and therefore
|
||||
# do not need to have a recomputation during backward.
|
||||
segment_ops = [self.program_ops[idx] for idx in segment]
|
||||
segment_len = len(segment)
|
||||
for idx in range(segment_len - 1, 0, -1):
|
||||
op = segment_ops[idx]
|
||||
user_ops = set()
|
||||
for res in op.results():
|
||||
user_ops = user_ops | set(res.all_used_ops())
|
||||
|
||||
if user_ops & set(segment_ops):
|
||||
continue
|
||||
segment.pop(idx)
|
||||
logger.info(
|
||||
f"Remove outgoing OP '{op.name()}' from the segment for recomputation, as it does not participate in the backward."
|
||||
)
|
||||
return segment
|
||||
|
||||
def get_segments(self):
|
||||
# `fwd_recompute_id` indicates the ID assigned to the segment for
|
||||
# which the OP requires recompute.
|
||||
# A segment comprises all OPs within a program, ranging from the OP
|
||||
# with the minimum index to the OP with the maximum index, and all
|
||||
# these operations share the same `fwd_recompute_id`.
|
||||
segment_beg = {}
|
||||
segment_end = {}
|
||||
max_op_id = len(self.program_ops)
|
||||
for idx, op in enumerate(self.program_ops):
|
||||
# 1. Find the OPs marked with `fwd_recompute_id`.
|
||||
if not op.has_attr("fwd_recompute_id"):
|
||||
continue
|
||||
# 2. Delineate the segment range marked by `fwd_recompute_id`.
|
||||
# Note: there may be some unmarked OPs in between.
|
||||
rc_id = op.attrs()["fwd_recompute_id"]
|
||||
if rc_id not in segment_beg:
|
||||
segment_beg[rc_id] = max_op_id
|
||||
segment_end[rc_id] = 0
|
||||
segment_beg[rc_id] = min(segment_beg[rc_id], idx)
|
||||
segment_end[rc_id] = max(segment_end[rc_id], idx)
|
||||
|
||||
# 3. Aggregate all segment information into a dictionary.
|
||||
# The key is the id of the segment, which is used to uniquely identify each segment.
|
||||
# The value is a list of indices of the segment OPs in `self.program_ops`.
|
||||
segments = {}
|
||||
assert len(segment_beg.keys()) == len(segment_end.keys())
|
||||
for segment_id, beg_id in segment_beg.items():
|
||||
assert segment_id in segment_end.keys()
|
||||
end_id = segment_end[segment_id]
|
||||
assert beg_id <= end_id
|
||||
segment = list(range(beg_id, end_id + 1))
|
||||
# 4. Remove the outgoing OPs from the segment, as these OPs
|
||||
# do not participate in the backward gradient computation.
|
||||
segments[segment_id] = self.remove_outgoing_op(segment)
|
||||
logger.info(
|
||||
f"Segment ID {segment_id} contains {len(segment)} OPs, all of which will be recomputed."
|
||||
)
|
||||
return segments
|
||||
|
||||
def get_op_name(self, op):
|
||||
return op.name().split('.')[1]
|
||||
|
||||
def match_pattern(
|
||||
self,
|
||||
op,
|
||||
visit,
|
||||
fetch_id,
|
||||
fetch_pattern,
|
||||
target_pattern,
|
||||
pre_len,
|
||||
main_len,
|
||||
count,
|
||||
max_count,
|
||||
):
|
||||
if count >= max_count:
|
||||
return max_count
|
||||
if len(fetch_pattern) > len(target_pattern):
|
||||
return count
|
||||
if self.get_op_name(op) != target_pattern[fetch_id]:
|
||||
return count
|
||||
if fetch_id == len(target_pattern) - 1:
|
||||
for idx in range(pre_len, pre_len + main_len):
|
||||
fetch_op = fetch_pattern[idx]
|
||||
visit[fetch_op] = -1
|
||||
refined_segment = list(set(visit.values()))
|
||||
refined_segment.sort()
|
||||
refined_segment = [idx for idx in refined_segment if idx != -1]
|
||||
return count + 1
|
||||
for res_val in op.results():
|
||||
for user_op in res_val.all_used_ops():
|
||||
fetch_pattern[fetch_id + 1] = user_op
|
||||
count = self.match_pattern(
|
||||
op=user_op,
|
||||
visit=visit,
|
||||
fetch_id=fetch_id + 1,
|
||||
fetch_pattern=fetch_pattern,
|
||||
target_pattern=target_pattern,
|
||||
pre_len=pre_len,
|
||||
main_len=main_len,
|
||||
count=count,
|
||||
max_count=max_count,
|
||||
)
|
||||
return count
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context=None):
|
||||
self.program_ops = list(main_program.global_block().ops)
|
||||
# 1. Get the recompute segments information form program.
|
||||
segments = self.get_segments()
|
||||
assert len(segments) > 0, (
|
||||
"No segment found in the PIR recompute pass.\n \
|
||||
Please disable 'recompute.enable' or check 'recompute()' usage in model code."
|
||||
)
|
||||
|
||||
# 2. Get the forward and backward OPs from program.
|
||||
fwd_ops, bwd_ops = self.get_fwd_bwd_ops()
|
||||
|
||||
# 3. Refine the segments based on the patterns.
|
||||
refined_ops_patterns = self.get_attr("refined_ops_patterns")
|
||||
for refined_ops_pattern in refined_ops_patterns:
|
||||
# 3.1 get the refined pattern information.
|
||||
# refined_ops_patterns = pre_ops + main_ops + suf_ops
|
||||
# `main_ops` pattern: it does not participate in backward recomputation
|
||||
# and needs to be removed from the segment.
|
||||
# `pre_ops` pattern: it serve only as markers and do require recomputation.
|
||||
# `suf_ops` pattern: it serve only as markers and do require recomputation.
|
||||
# `num` : it limits the maximum number of `main_ops` patterns identified
|
||||
# within each segment. A value of -1 represents all patterns.
|
||||
num = int(refined_ops_pattern['num'])
|
||||
num = num if num >= 0 else len(fwd_ops)
|
||||
main_ops = refined_ops_pattern['main_ops']
|
||||
pre_ops = refined_ops_pattern['pre_ops']
|
||||
suf_ops = refined_ops_pattern['suf_ops']
|
||||
pattern_ops = pre_ops + main_ops + suf_ops
|
||||
|
||||
for rc_id in segments.keys():
|
||||
# 3.2 Identify and mark the first 'num' patterns in each segment.
|
||||
# The dictionary 'op_idx_map' has keys as OP information.
|
||||
# If an OP belongs to a pattern, its value in the dictionary is marked as -1.
|
||||
op_idx_map = {
|
||||
self.program_ops[idx]: idx for idx in segments[rc_id]
|
||||
}
|
||||
pattern_count = 0
|
||||
fetch_pattern = [None] * len(pattern_ops)
|
||||
for idx in segments[rc_id]:
|
||||
op = self.program_ops[idx]
|
||||
fetch_pattern[0] = op
|
||||
pattern_count = self.match_pattern(
|
||||
op=self.program_ops[idx],
|
||||
visit=op_idx_map,
|
||||
fetch_id=0,
|
||||
fetch_pattern=fetch_pattern,
|
||||
target_pattern=pattern_ops,
|
||||
pre_len=len(pre_ops),
|
||||
main_len=len(main_ops),
|
||||
count=pattern_count,
|
||||
max_count=num,
|
||||
)
|
||||
# 3.3 Refined segment to exclude the specified pattern.
|
||||
refined_segment = list(set(op_idx_map.values()))
|
||||
refined_segment.sort()
|
||||
refined_segment = [idx for idx in refined_segment if idx != -1]
|
||||
segments[rc_id] = refined_segment
|
||||
|
||||
# 4. Construct the segment for backward recomputation.
|
||||
# 4.1 Build IrMapping to eplace forward value with backward recompute value.
|
||||
input_value = main_program.list_vars()
|
||||
value_map = paddle.pir.IrMapping()
|
||||
for val in input_value:
|
||||
value_map.add(val, val)
|
||||
|
||||
for rc_id, segment in segments.items():
|
||||
# 4.2 Find the insertion position for the backward segment,
|
||||
# which should be before backward gradient computation.
|
||||
first_bwd_used_op = bwd_ops[-1]
|
||||
for idx in segment:
|
||||
op = self.program_ops[idx]
|
||||
bwd_used_op = self.get_first_bwd_used_op(op, bwd_ops)
|
||||
if first_bwd_used_op.id() > bwd_used_op.id():
|
||||
first_bwd_used_op = bwd_used_op
|
||||
|
||||
ori_segment_outputs = backward_utils.ValueSet()
|
||||
paddle.pir.set_insertion_point(first_bwd_used_op)
|
||||
|
||||
# 4.3 Clone the segment OPs and replace the forward
|
||||
# value with backward recompute value.
|
||||
for idx in segment:
|
||||
op = self.program_ops[idx]
|
||||
ori_segment_outputs.update(op.results())
|
||||
|
||||
# Random OPs should produce the same output before and after recomputation.
|
||||
if self.is_seed_used_by_dropout(op):
|
||||
continue
|
||||
|
||||
rc_op = op.clone(
|
||||
value_map, paddle.pir.CloneOptions(False, True, True)
|
||||
)
|
||||
# The forward segment and the backward segment have the same segment ID.
|
||||
if rc_op.has_attr("fwd_recompute_id"):
|
||||
rc_op.erase_attr("fwd_recompute_id")
|
||||
|
||||
rc_op.set_int_attr("bwd_recompute_id", rc_id)
|
||||
|
||||
# Updtate attributes.
|
||||
if first_bwd_used_op.has_attr('op_role'):
|
||||
rc_op.set_int_attr("op_role", first_bwd_used_op.op_role)
|
||||
|
||||
if first_bwd_used_op.has_attr('chunk_id'):
|
||||
rc_op.set_int_attr("chunk_id", first_bwd_used_op.chunk_id)
|
||||
|
||||
# 4.4 Replace the forward value with backward recompute value.
|
||||
for ori_value in ori_segment_outputs:
|
||||
rc_value = value_map.look_up(ori_value)
|
||||
ori_value.replace_grad_users_with(rc_value, set(bwd_ops))
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
# 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
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.auto_parallel.static.process_group import (
|
||||
new_process_group,
|
||||
)
|
||||
|
||||
from ..auto_parallel.static.utils import (
|
||||
get_logger,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
@register_pass("replace_with_parallel_cross_entropy")
|
||||
class AutoParallelReplaceWithParallelCrossEntropyPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
hcg = dist.fleet.get_hybrid_communicate_group()
|
||||
self.model_parallel_group = hcg.get_model_parallel_group()
|
||||
self.tensor_parallel_degree = hcg.get_model_parallel_world_size()
|
||||
|
||||
def _check_self(self):
|
||||
# The activation of this pass requires adopting a model parallel strategy.
|
||||
if self.tensor_parallel_degree < 2:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _check_user(self, value):
|
||||
placement1 = value.placements
|
||||
for user in value.all_used_ops():
|
||||
for operand in user.operands_source():
|
||||
if operand.get_defining_op() != value.get_defining_op():
|
||||
continue
|
||||
placement2 = operand.placements
|
||||
if placement1 != placement2:
|
||||
return False
|
||||
break
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
del_ops = []
|
||||
new_ops = []
|
||||
|
||||
for block in main_program.blocks:
|
||||
for op in reversed(block.ops):
|
||||
if op.name() == 'pd_op.cross_entropy_with_softmax':
|
||||
operand1 = op.operand_source(0)
|
||||
operand2 = op.operand_source(1)
|
||||
|
||||
# The `logit` input of the `cross_stropy_with_stoftmax` operator
|
||||
# meed split along the column.
|
||||
placement1 = operand1.placements
|
||||
if not placement1[1].is_shard():
|
||||
return
|
||||
|
||||
process_ids = operand1.dist_attr().process_mesh.process_ids
|
||||
group = new_process_group(sorted(process_ids))
|
||||
ring_id = group.id
|
||||
nranks = group.nranks
|
||||
rank = paddle.distributed.get_rank()
|
||||
|
||||
ignore_index = op.attrs()["ignore_index"]
|
||||
paddle.pir.set_insertion_point(op)
|
||||
softmax, loss = paddle._C_ops.c_softmax_with_cross_entropy(
|
||||
operand1, operand2, ignore_index, ring_id, rank, nranks
|
||||
)
|
||||
op.result(0).replace_all_uses_with(softmax)
|
||||
op.result(1).replace_all_uses_with(loss)
|
||||
del_ops.append(op)
|
||||
new_ops.append(softmax.get_defining_op())
|
||||
|
||||
for op in del_ops:
|
||||
for result in op.results():
|
||||
assert result.use_empty()
|
||||
op.erase()
|
||||
# In the forward program, the placements of the newly added OP
|
||||
# output should be consistent with the placements of the user OP input
|
||||
for op in new_ops:
|
||||
for result in op.results():
|
||||
assert self._check_user(result)
|
||||
return
|
||||
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
naive_set_dist_op_attr_for_program_by_mesh,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY
|
||||
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
||||
from paddle.static import default_main_program
|
||||
|
||||
from .auto_parallel_sharding import _is_reshard_op
|
||||
from .pass_base import PassBase, PassType, register_pass
|
||||
|
||||
|
||||
# NOTE we add the "auto_parallel" prefix to the pass in order to
|
||||
# indicate that this pass should obey some constrains by auto_parallel
|
||||
# for example all ops and vars should has dist attr before and after pass
|
||||
# should use dist op instead of custom comm op
|
||||
@register_pass("auto_parallel_sequence_parallel_optimization")
|
||||
class SequenceParallelOptimizationPass(PassBase):
|
||||
"""
|
||||
This pass is used to optimize the sequence parallel.
|
||||
1. Fuse the allreduce + split into reducescatter.
|
||||
2. Trade off communication for memory in the row_parallel_linear output.
|
||||
3. Overlap communication with computation in backward computation.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("dist_context", None)
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
if (not isinstance(self.get_attr("global_rank"), int)) or self.get_attr(
|
||||
"global_rank"
|
||||
) < 0:
|
||||
return False
|
||||
if not self.get_attr("dist_context").strategy.sp_optimization.enable:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _type(self):
|
||||
return PassType.COMM_OPT
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
self.dist_context = self.get_attr("dist_context")
|
||||
self.global_rank = int(self.get_attr("global_rank"))
|
||||
|
||||
with paddle.static.program_guard(main_program, startup_program):
|
||||
# TODO remove this pass when we use local reshard for all communication
|
||||
self._fuse_allreduce_split()
|
||||
self._memory_optimization()
|
||||
self._overlap()
|
||||
|
||||
def _fuse_allreduce_split(self):
|
||||
# allreduce is added by dist op and split is added by reshard, so we need this pass to fuse them as reducescatter.
|
||||
# reducescatter should be inferred by local reshard in future.
|
||||
|
||||
block = default_main_program().global_block()
|
||||
|
||||
# record valid split ops
|
||||
valid_split_op_indices = []
|
||||
|
||||
def is_valid_split_op(idx, block):
|
||||
op = block.ops[idx]
|
||||
if not op.type == "split":
|
||||
return False
|
||||
pre_op = block.ops[idx - 1]
|
||||
if not (
|
||||
pre_op.type == "all_reduce"
|
||||
and pre_op.attr("reduce_type")
|
||||
== paddle.distributed.ReduceOp.SUM
|
||||
):
|
||||
return False
|
||||
pre_output_name = pre_op.output_arg_names[0]
|
||||
cur_input_name = op.input_arg_names[0]
|
||||
if (
|
||||
pre_output_name == cur_input_name
|
||||
and _is_reshard_op(op)
|
||||
and op.attr("axis") == 0
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
for i in range(len(block.ops)):
|
||||
if is_valid_split_op(i, block):
|
||||
valid_split_op_indices.append(i)
|
||||
|
||||
# modify program
|
||||
remove_varnames = []
|
||||
for i in sorted(valid_split_op_indices, reverse=True):
|
||||
allreduce_op = block.ops[i - 1]
|
||||
split_op = block.ops[i]
|
||||
consumer_op = block.ops[i + 1]
|
||||
|
||||
allreduce_input_name = allreduce_op.input("X")[0]
|
||||
ring_id = int(allreduce_op.attr("ring_id"))
|
||||
split_output_names = split_op.output("Out")
|
||||
nranks = len(split_output_names)
|
||||
consumer_input_names = consumer_op.input_arg_names
|
||||
intersection = set(split_output_names).intersection(
|
||||
set(consumer_input_names)
|
||||
)
|
||||
assert len(intersection) == 1, (
|
||||
f"Sequence Parallel ReduceScatter Output more than 1: {intersection}."
|
||||
)
|
||||
keep_output_name = intersection.pop()
|
||||
split_output_names.remove(keep_output_name)
|
||||
remove_varnames.extend(split_output_names)
|
||||
|
||||
# replace ops
|
||||
new_op = block._insert_op_without_sync(
|
||||
index=i + 1,
|
||||
type="reduce_scatter",
|
||||
inputs={'x': [allreduce_input_name]},
|
||||
outputs={'out': [keep_output_name]},
|
||||
attrs={
|
||||
'ring_id': ring_id,
|
||||
'nranks': nranks,
|
||||
'op_namescope': allreduce_op.attr("op_namescope"),
|
||||
OP_ROLE_KEY: consumer_op.attr(OP_ROLE_KEY),
|
||||
},
|
||||
)
|
||||
new_op.dist_attr.execution_stream = (
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
block._remove_op(i, False)
|
||||
block._remove_op(i - 1, False)
|
||||
|
||||
# set dist attr
|
||||
allreduce_input_dist_attr = (
|
||||
self.dist_context.get_tensor_dist_attr_for_program(
|
||||
block.vars[allreduce_input_name]
|
||||
)
|
||||
)
|
||||
ref_process_mesh = allreduce_input_dist_attr.process_mesh
|
||||
naive_set_dist_op_attr_for_program_by_mesh(
|
||||
new_op,
|
||||
ref_process_mesh,
|
||||
self.dist_context,
|
||||
chunk_id=allreduce_input_dist_attr.chunk_id,
|
||||
)
|
||||
|
||||
# remove vars
|
||||
for varname in remove_varnames:
|
||||
block._remove_var(varname, sync=False)
|
||||
|
||||
block._sync_with_cpp()
|
||||
|
||||
def _memory_optimization(self):
|
||||
pass
|
||||
|
||||
def _overlap(self):
|
||||
pass
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,159 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
from paddle.distributed.auto_parallel.static.operators.common import (
|
||||
is_amp_flag_sync_op,
|
||||
is_data_parallel_reduce_op,
|
||||
is_global_norm_sync_op,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
OpRole,
|
||||
insert_dependencies_for_vars,
|
||||
is_comm_op,
|
||||
)
|
||||
|
||||
from .auto_parallel_sharding import ShardingPass, _supported_optimizer_type
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
|
||||
def _sharding_pass_applied(pass_ctx):
|
||||
for applied_pass in pass_ctx.passes:
|
||||
if isinstance(applied_pass, ShardingPass):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# NOTE we add the "auto_parallel" prefix to the pass in order to
|
||||
# indicate that this pass should obey some constrains by auto_parallel
|
||||
# for example all ops and vars should has dist attr before and after pass
|
||||
# should use dist op instead of custom comm op
|
||||
@register_pass("auto_parallel_supplement_explicit_dependencies")
|
||||
class AutoParalSupplementDepPass(PassBase):
|
||||
"""
|
||||
Functional Concern.
|
||||
for strategies like amp & global norm, there is a collective communication to sync gradient information in every rank.
|
||||
after partition the gradients to each rank, the order of that collective communication is different in each rank
|
||||
and might cause hang problem in graph based random order executor. here supplement explicit dependencies for those cases.
|
||||
|
||||
TODO Performance Concern.
|
||||
global collective will introduce global synchronization which forces the fast workers to wait for slow ones.
|
||||
therefore we should conduct this collective when all the ranks reach a same stage.
|
||||
BUT the depend API offered by executor could only ensure "conduct-not-before" but not "conduct-right-after".
|
||||
Some ranks might call the collectives first than other ranks while they still some local could be performed to wait for slow peers.
|
||||
IR Pass currently could not have the fully control of time the to perform these global collectives.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("dist_context", None)
|
||||
|
||||
def _check_self(self):
|
||||
if self.get_attr("dist_context") is None:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
# TODO general this pass for all case.
|
||||
if not _sharding_pass_applied(context):
|
||||
return
|
||||
|
||||
self._dist_context = self.get_attr("dist_context", None)
|
||||
self.flags_sync_stream = "flags_sync_stream"
|
||||
main_block = main_program.global_block()
|
||||
startup_block = startup_program.global_block()
|
||||
|
||||
# last dp grad communication
|
||||
last_dp_reduce_op_idx = -1
|
||||
last_dp_reduce_varname = None
|
||||
for idx, op in reversed(list(enumerate(main_block.ops))):
|
||||
if is_data_parallel_reduce_op(op):
|
||||
last_dp_reduce_op_idx = idx
|
||||
last_dp_reduce_varname = op.output_arg_names[0]
|
||||
break
|
||||
assert last_dp_reduce_op_idx > 0
|
||||
assert last_dp_reduce_varname is not None
|
||||
|
||||
# analyze deps for amp & global norm
|
||||
deps_map = {}
|
||||
prior_varname = last_dp_reduce_varname
|
||||
for idx, op in enumerate(main_block.ops):
|
||||
if is_amp_flag_sync_op(op) or is_global_norm_sync_op(op):
|
||||
op_namescope = None
|
||||
if is_amp_flag_sync_op(op):
|
||||
op_namescope = "amp_flag_sync_dep"
|
||||
op.dist_attr.execution_stream = self.flags_sync_stream
|
||||
|
||||
elif is_global_norm_sync_op(op):
|
||||
op_namescope = "global_norm_sync_dep"
|
||||
deps_map[idx] = (prior_varname, op.input("X")[0], op_namescope)
|
||||
prior_varname = op.output("Out")[0]
|
||||
|
||||
# analyze deps for check_finite_and_unscale
|
||||
# ensure it is performed after last backward computation, therefore reduce the
|
||||
# straggling of the amp-flag-sync
|
||||
first_check_op = True
|
||||
for idx, op in enumerate(main_block.ops):
|
||||
if op.type == "check_finite_and_unscale":
|
||||
if first_check_op:
|
||||
last_backward_op = None
|
||||
for last_idx in range(idx - 1, 0, -1):
|
||||
if not is_comm_op(main_block.ops[last_idx]):
|
||||
last_backward_op = main_block.ops[last_idx]
|
||||
break
|
||||
prior_varname = last_backward_op.output_arg_names[0]
|
||||
first_check_op = False
|
||||
deps_map[idx] = (
|
||||
prior_varname,
|
||||
op.input("Scale")[0],
|
||||
"check_finite_dep",
|
||||
)
|
||||
|
||||
# analyze deps for optimizer
|
||||
# optimizers order should be fixed to allow broadcast to overlap with optimizer
|
||||
first_optimizer_op = True
|
||||
for idx, op in enumerate(main_block.ops):
|
||||
if op.type in _supported_optimizer_type:
|
||||
if first_optimizer_op:
|
||||
first_optimizer_op = False
|
||||
else:
|
||||
deps_map[idx] = (
|
||||
prior_varname,
|
||||
op.input("Param")[0],
|
||||
"optimizer_order_dep",
|
||||
)
|
||||
prior_varname = op.output("ParamOut")[0]
|
||||
|
||||
# insert deps
|
||||
indice = sorted(deps_map.keys(), reverse=True)
|
||||
for idx in indice:
|
||||
prior_var = main_block.var(deps_map[idx][0])
|
||||
post_var = main_block.var(deps_map[idx][1])
|
||||
op_namescope = deps_map[idx][2]
|
||||
depend_op = insert_dependencies_for_vars(
|
||||
main_block,
|
||||
idx,
|
||||
prior_var,
|
||||
post_var,
|
||||
self._dist_context,
|
||||
OpRole.Optimize,
|
||||
is_recompute=False,
|
||||
sync=False,
|
||||
op_namescope=op_namescope,
|
||||
)
|
||||
|
||||
main_block._sync_with_cpp()
|
||||
@@ -0,0 +1,302 @@
|
||||
# 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
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.base.core import TensorDistAttr
|
||||
from paddle.base.executor import global_scope
|
||||
from paddle.base.framework import auto_complete_op_role
|
||||
from paddle.distributed.auto_parallel.static.process_group import (
|
||||
new_process_group,
|
||||
)
|
||||
from paddle.distributed.auto_parallel.static.utils import (
|
||||
get_pp_stage_by_process_mesh,
|
||||
)
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
from paddle.static.pir_io import get_pir_parameters
|
||||
|
||||
from ..auto_parallel.static.utils import (
|
||||
get_logger,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
@register_pass("auto_parallel_sync_shared_params")
|
||||
class AutoParallelSyncSharedParamsPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.params_maybe_shared = []
|
||||
self.src_ranks = []
|
||||
self.dst_ranks = []
|
||||
self.comm_group = {}
|
||||
|
||||
def _check_self(self):
|
||||
pipeline_strategy = self.get_attr('pipeline_strategy')
|
||||
if (not pipeline_strategy.enable) or pipeline_strategy.pp_degree <= 1:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _find_fist_opt_user(self, main_program):
|
||||
for op in main_program.global_block().ops:
|
||||
if op.op_role == 2:
|
||||
return op
|
||||
|
||||
def _get_comm_group(self, ranks=[]):
|
||||
ranks = sorted(ranks)
|
||||
if tuple(ranks) in self.comm_group:
|
||||
return self.comm_group[tuple(ranks)]
|
||||
# The communication group of this `all_reduce` op satisfies len (ranks)==2.
|
||||
# When `force_new_group=False` is set, the `send&recv` group will be returned,
|
||||
# At this point, `all_reduce` and `send&recv` share the same group, and
|
||||
# the process will hang up.
|
||||
group = new_process_group(ranks, force_new_group=True)
|
||||
self.comm_group[tuple(ranks)] = group.id
|
||||
return group.id
|
||||
|
||||
def sync_shared_parameters(self, main_program, startup_program):
|
||||
if not self._check_self():
|
||||
logger.info(
|
||||
"AutoParallelSyncSharedParamsPass need support pipeline parallel, skip pass."
|
||||
)
|
||||
return []
|
||||
new_shared_params = []
|
||||
params, _ = get_pir_parameters(main_program)
|
||||
for param in params:
|
||||
users = param.all_used_ops()
|
||||
for user_op in users:
|
||||
if user_op.name() == "dist_op.reshard":
|
||||
reshard_op = user_op
|
||||
dist_attr = reshard_op.dist_attr
|
||||
src_dist_attr = dist_attr.operand(0).as_tensor_dist_attr()
|
||||
dst_dist_attr = dist_attr.result(0).as_tensor_dist_attr()
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
# Shared parameter needs reshard on diff stage.
|
||||
pipeline_strategy = self.get_attr('pipeline_strategy')
|
||||
pp_degree = pipeline_strategy.pp_degree
|
||||
src_stage = get_pp_stage_by_process_mesh(
|
||||
src_mesh, pp_degree
|
||||
)
|
||||
dst_stage = get_pp_stage_by_process_mesh(
|
||||
dst_mesh, pp_degree
|
||||
)
|
||||
if (
|
||||
src_stage is None
|
||||
or dst_stage is None
|
||||
or src_stage == dst_stage
|
||||
):
|
||||
continue
|
||||
|
||||
# Get shared parameter name
|
||||
param_name = param.get_defining_op().str_attr(
|
||||
'parameter_name'
|
||||
)
|
||||
|
||||
# Add shared parameter builtin.parameter with "shared_" prefix.
|
||||
with (
|
||||
auto_complete_op_role(main_program, OpRole.Forward),
|
||||
paddle.static.program_guard(
|
||||
main_program, startup_program
|
||||
),
|
||||
):
|
||||
shared_param = paddle.pir.core.create_parameter(
|
||||
dtype=param.dtype,
|
||||
shape=param.shape,
|
||||
name="shared_" + param_name,
|
||||
process_mesh=dst_mesh,
|
||||
placements=src_dist_attr.placements,
|
||||
initializer=paddle.nn.initializer.Constant(value=0),
|
||||
)
|
||||
main_program.set_parameters_from(startup_program)
|
||||
|
||||
# Record new shared parameter.
|
||||
new_shared_params.append("shared_" + param_name)
|
||||
|
||||
# Set value for new shared parameter.
|
||||
concrete_program = self.get_attr("concrete_program")
|
||||
dy_params = concrete_program.parameters[0]
|
||||
dy_param = None
|
||||
for tmp_param in dy_params:
|
||||
if tmp_param.name == param_name:
|
||||
dy_param = tmp_param
|
||||
break
|
||||
assert dy_param is not None, (
|
||||
f"The parameter {param_name} was not found in the concrete_degram"
|
||||
)
|
||||
|
||||
new_dist_attr = TensorDistAttr()
|
||||
new_dist_attr.process_mesh = dst_mesh
|
||||
new_dist_attr.dims_mapping = src_dist_attr.dims_mapping
|
||||
with paddle.no_grad():
|
||||
dy_shared_param = paddle.base.core.reshard(
|
||||
dy_param, new_dist_attr
|
||||
)
|
||||
paddle.device.synchronize()
|
||||
if dy_shared_param._is_initialized():
|
||||
pir_shared_param = (
|
||||
global_scope()
|
||||
.var("shared_" + param_name)
|
||||
.get_tensor()
|
||||
)
|
||||
pir_shared_param._share_data_with(
|
||||
dy_shared_param.get_tensor().get_tensor()
|
||||
)
|
||||
|
||||
# record in params_maybe_shared
|
||||
self.params_maybe_shared.append(
|
||||
{
|
||||
'src_mesh': src_mesh,
|
||||
'dst_mesh': dst_mesh,
|
||||
'src_dist_attr': src_dist_attr,
|
||||
'dst_dist_attr': dst_dist_attr,
|
||||
'param_name': param_name,
|
||||
}
|
||||
)
|
||||
|
||||
# New shared parameter must has same dist_attr with shared parameter
|
||||
new_src_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
)
|
||||
if new_src_dist_attr == dst_dist_attr:
|
||||
# Remove useless reshared op.
|
||||
reshard_op.result(0).replace_all_uses_with(shared_param)
|
||||
reshard_op.erase()
|
||||
|
||||
else:
|
||||
# Update reshard op dist_attr.
|
||||
reshard_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh,
|
||||
[new_src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
-1,
|
||||
)
|
||||
)
|
||||
reshard_op.operand(0).set_source(shared_param)
|
||||
|
||||
self.src_ranks.extend(src_mesh.process_ids)
|
||||
self.dst_ranks.extend(dst_mesh.process_ids)
|
||||
|
||||
if len(self.params_maybe_shared) == 0:
|
||||
logger.info("No parameter need to share, skip pass.")
|
||||
return []
|
||||
|
||||
# Must initialize the redundant communication group for the allreduce op here.
|
||||
# Otherwise, it will hang during gradient synchronization.
|
||||
for idx in range(len(self.src_ranks)):
|
||||
rank_1 = self.src_ranks[idx]
|
||||
rank_2 = self.dst_ranks[idx]
|
||||
new_process_group(sorted([rank_1, rank_2]))
|
||||
self._get_comm_group([rank_1, rank_2])
|
||||
|
||||
return new_shared_params
|
||||
|
||||
def sync_shared_parameter_gradient(
|
||||
self, main_program, startup_program, params_grads
|
||||
):
|
||||
if not self._check_self():
|
||||
logger.info(
|
||||
"AutoParallelSyncSharedParamsPass need support pipeline parallel, skip pass."
|
||||
)
|
||||
return params_grads
|
||||
|
||||
if len(self.params_maybe_shared) == 0:
|
||||
logger.info("No parameter need to share, skip pass.")
|
||||
return params_grads
|
||||
|
||||
# Only support one shared parameter.
|
||||
# TODO: support more shared parameters
|
||||
assert len(self.params_maybe_shared) == 1, (
|
||||
"Currently, only one shared parameter is supported, and it cannot support more at the moment."
|
||||
)
|
||||
|
||||
cur_rank = paddle.distributed.get_rank()
|
||||
|
||||
if cur_rank not in self.src_ranks and cur_rank not in self.dst_ranks:
|
||||
return params_grads
|
||||
|
||||
pre_name = ""
|
||||
if cur_rank in self.dst_ranks:
|
||||
pre_name = "shared_"
|
||||
|
||||
for param_mess in self.params_maybe_shared:
|
||||
param_name = pre_name + param_mess['param_name']
|
||||
src_mesh_ids = param_mess['src_mesh'].process_ids
|
||||
dst_mesh_ids = param_mess['dst_mesh'].process_ids
|
||||
|
||||
# Get (param, grad) value
|
||||
param_value = main_program.get_parameter_value_by_name(param_name)
|
||||
|
||||
grad_idx = None
|
||||
for p_idx, (p_param, _) in enumerate(params_grads):
|
||||
if p_param.is_same(param_value):
|
||||
grad_idx = p_idx
|
||||
break
|
||||
assert grad_idx is not None, (
|
||||
f"Parameter {param_name} not found in params_grades, unable to find corresponding gradient value."
|
||||
)
|
||||
grad_value = params_grads[p_idx][1]
|
||||
|
||||
# Create allreduce op comm group.
|
||||
cur_rank = paddle.distributed.get_rank()
|
||||
if cur_rank in self.src_ranks:
|
||||
idx = src_mesh_ids.index(cur_rank)
|
||||
peer_rank = dst_mesh_ids[idx]
|
||||
if cur_rank in self.dst_ranks:
|
||||
idx = dst_mesh_ids.index(cur_rank)
|
||||
peer_rank = src_mesh_ids[idx]
|
||||
ar_group_id = self._get_comm_group([cur_rank, peer_rank])
|
||||
|
||||
# Insert allreduce op in the end of backward.
|
||||
insert_pos = self._find_fist_opt_user(main_program)
|
||||
paddle.pir.set_insertion_point(insert_pos)
|
||||
|
||||
# Build allreduce op to sync gradient.
|
||||
with auto_complete_op_role(main_program, OpRole.Backward):
|
||||
allreduce_val = paddle._C_ops.all_reduce(
|
||||
grad_value,
|
||||
ar_group_id,
|
||||
dist.ReduceOp.SUM,
|
||||
)
|
||||
allreduce_val.update_dist_attr(grad_value.dist_attr())
|
||||
allreduce_op = allreduce_val.get_defining_op()
|
||||
|
||||
# Update all_used_ops
|
||||
for user in grad_value.all_used_ops():
|
||||
if user.name() == "pd_op.all_reduce":
|
||||
continue
|
||||
for idx, operand in enumerate(user.operands()):
|
||||
if user.operand_source(idx).is_same(grad_value):
|
||||
user.operand(idx).set_source(allreduce_val)
|
||||
|
||||
# Update (param, grad) value
|
||||
params_grads[p_idx] = (param_value, allreduce_val)
|
||||
|
||||
return params_grads
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
return
|
||||
Executable
+251
@@ -0,0 +1,251 @@
|
||||
# Copyright (c) 2021 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.
|
||||
|
||||
from paddle.framework import (
|
||||
_apply_pass as _apply_cpp_pass,
|
||||
core,
|
||||
)
|
||||
from paddle.static import Executor
|
||||
|
||||
from .pass_base import CPPPassWrapper, PassType, register_pass
|
||||
|
||||
|
||||
@register_pass("fuse_elewise_add_act")
|
||||
class FuseElementwiseAddActPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_elewise_add_act_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_bn_act")
|
||||
class FuseBatchNormActPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_bn_act_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_bn_add_act")
|
||||
class FuseBatchNormAddActPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_bn_add_act_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_relu_depthwise_conv")
|
||||
class FuseReluDepthwiseConvPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_relu_depthwise_conv_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fused_attention")
|
||||
class FusedAttentionPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fused_attention_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fused_feedforward")
|
||||
class FusedFeedforwardPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fused_feedforward_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_gemm_epilogue")
|
||||
class FuseGemmEpiloguePass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_gemm_epilogue_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_adamw")
|
||||
class FuseAdamWPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_adamw_op_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_dot_product_attention")
|
||||
class FuseDotProductAttentionPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_dot_product_attention_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_optimizer")
|
||||
class FuseOptimizerPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return [
|
||||
"fuse_adam_op_pass",
|
||||
"fuse_sgd_op_pass",
|
||||
"fuse_momentum_op_pass",
|
||||
]
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
# @register_pass("inplace_addto_op")
|
||||
# class InplaceAddtoOpPass(CPPPassWrapper):
|
||||
# def __init__(self):
|
||||
# super().__init__()
|
||||
|
||||
# @property
|
||||
# def cpp_name(self):
|
||||
# return "inplace_addto_op_pass"
|
||||
|
||||
# def _type(self):
|
||||
# return PassType.CALC_OPT
|
||||
|
||||
|
||||
@register_pass("fuse_resunit")
|
||||
class FuseResUnitPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "fuse_resunit_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.FUSION_OPT
|
||||
|
||||
|
||||
def _set_cinn_op_flag(flag_name, extra_ops):
|
||||
values = core.globals()[flag_name]
|
||||
values = [v.strip() for v in values.split(";") if v.strip()]
|
||||
values.extend(extra_ops)
|
||||
core.globals()[flag_name] = ";".join(values)
|
||||
|
||||
|
||||
@register_pass("build_cinn")
|
||||
class BuildCINNPass(CPPPassWrapper):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("allow_ops", [])
|
||||
self.set_attr("deny_ops", [])
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
return "build_cinn_pass"
|
||||
|
||||
def _type(self):
|
||||
return PassType.CALC_OPT
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
assert 'FLAGS_allow_cinn_ops' in core.globals(), (
|
||||
"PaddlePaddle is not compiled with CINN support"
|
||||
)
|
||||
old_allow_ops = core.globals()['FLAGS_allow_cinn_ops']
|
||||
old_deny_ops = core.globals()['FLAGS_deny_cinn_ops']
|
||||
try:
|
||||
_set_cinn_op_flag(
|
||||
'FLAGS_allow_cinn_ops', self.get_attr("allow_ops")
|
||||
)
|
||||
_set_cinn_op_flag('FLAGS_deny_cinn_ops', self.get_attr("deny_ops"))
|
||||
|
||||
feed = self.get_attr('feed', [])
|
||||
fetch_list = self.get_attr('fetch_list', [])
|
||||
prune_program = self.get_attr('prune_program', True)
|
||||
|
||||
if prune_program:
|
||||
tmp_main_program = Executor._prune_program(
|
||||
main_program, feed, fetch_list, []
|
||||
)
|
||||
|
||||
tmp_main_program = Executor._add_fetch_ops(
|
||||
tmp_main_program, fetch_list, 'fetch'
|
||||
)
|
||||
|
||||
else:
|
||||
tmp_main_program = Executor._add_fetch_ops(
|
||||
main_program, fetch_list, 'fetch'
|
||||
)
|
||||
|
||||
_apply_cpp_pass(
|
||||
tmp_main_program,
|
||||
startup_program,
|
||||
self.cpp_name,
|
||||
{},
|
||||
self.cpp_attr_types,
|
||||
)
|
||||
|
||||
tmp_main_program = Executor._remove_fetch_ops(tmp_main_program)
|
||||
|
||||
tmp_main_program = core.ProgramDesc(tmp_main_program.desc)
|
||||
|
||||
main_program._rebuild_from_desc(tmp_main_program)
|
||||
|
||||
finally:
|
||||
core.globals()['FLAGS_allow_cinn_ops'] = old_allow_ops
|
||||
core.globals()['FLAGS_deny_cinn_ops'] = old_deny_ops
|
||||
+387
@@ -0,0 +1,387 @@
|
||||
# Copyright (c) 2021 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 numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.framework import core
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from .pass_base import PassBase, PassType, register_pass
|
||||
|
||||
|
||||
def find_adjacent_match_sequences(
|
||||
iterable, filter_func, adjacent_filter_func=None
|
||||
):
|
||||
n = len(iterable)
|
||||
match_sequences = []
|
||||
if adjacent_filter_func is None:
|
||||
adjacent_filter_func = lambda ref_op, new_op: True
|
||||
i = 0
|
||||
while True:
|
||||
while i < n and not filter_func(iterable[i]):
|
||||
i += 1
|
||||
j = i + 1
|
||||
while (
|
||||
j < n
|
||||
and filter_func(iterable[j])
|
||||
and adjacent_filter_func(iterable[i], iterable[j])
|
||||
):
|
||||
j += 1
|
||||
if i < n and j <= n:
|
||||
match_sequences.append((i, j))
|
||||
i = j + 1
|
||||
if i >= n:
|
||||
break
|
||||
return match_sequences
|
||||
|
||||
|
||||
def insert_fuse_all_reduce_ops(
|
||||
block, reversed_op_indices, input_var_names, output_var_names, dtype, attrs
|
||||
):
|
||||
fused_var = block.create_var(
|
||||
name=unique_name.generate(f"FusedOutput_{input_var_names[0]}"),
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# FIXME(zengjinle): here we assume that we use
|
||||
# c_sync_calc_stream/c_sync_comm_stream to do sync.
|
||||
# But someone may use c_wait_compute/c_wait_comm instead.
|
||||
if not attrs["use_calc_stream"]:
|
||||
ring_id = attrs["ring_id"]
|
||||
new_op_indices = list(reversed_op_indices)
|
||||
|
||||
for i, op_idx in enumerate(reversed_op_indices):
|
||||
prev_op_idx = op_idx - 1
|
||||
while (
|
||||
prev_op_idx >= 0
|
||||
and block.ops[prev_op_idx].type == "c_sync_calc_stream"
|
||||
):
|
||||
new_op_indices.append(prev_op_idx)
|
||||
prev_op_idx -= 1
|
||||
|
||||
if i > 0:
|
||||
next_op_idx = op_idx + 1
|
||||
n = len(block.ops)
|
||||
while (
|
||||
next_op_idx < n
|
||||
and block.ops[next_op_idx].type == "c_sync_comm_stream"
|
||||
):
|
||||
assert block.ops[next_op_idx].attr("ring_id") == ring_id
|
||||
new_op_indices.append(next_op_idx)
|
||||
|
||||
new_op_indices = list(set(new_op_indices))
|
||||
new_op_indices.sort(reverse=True)
|
||||
reversed_op_indices = new_op_indices
|
||||
|
||||
insert_idx = reversed_op_indices[0] + 1
|
||||
op_role_key = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
||||
|
||||
concated_shapes = []
|
||||
concated_ranks = []
|
||||
for var_name in output_var_names:
|
||||
shape = block._find_var_recursive(var_name).shape
|
||||
concated_shapes.extend(shape)
|
||||
concated_ranks.append(len(shape))
|
||||
|
||||
coalesce_tensor_op_kwargs = {
|
||||
"type": "coalesce_tensor",
|
||||
"inputs": {
|
||||
"Input": input_var_names,
|
||||
},
|
||||
"outputs": {
|
||||
"Output": output_var_names,
|
||||
"FusedOutput": fused_var,
|
||||
},
|
||||
"attrs": {
|
||||
"use_align": True,
|
||||
"dtype": dtype,
|
||||
"concated_shapes": concated_shapes,
|
||||
"concated_ranks": concated_ranks,
|
||||
op_role_key: attrs[op_role_key],
|
||||
},
|
||||
}
|
||||
|
||||
if not attrs["use_calc_stream"]:
|
||||
block._insert_op_without_sync(
|
||||
insert_idx,
|
||||
type="c_sync_calc_stream",
|
||||
inputs={"X": fused_var},
|
||||
outputs={"Out": fused_var, op_role_key: attrs[op_role_key]},
|
||||
)
|
||||
insert_idx += 1
|
||||
|
||||
# all_reduce sum should insert
|
||||
attrs["reduce_type"] = paddle.distributed.ReduceOp.SUM
|
||||
block._insert_op_without_sync(
|
||||
insert_idx,
|
||||
type="all_reduce",
|
||||
inputs={"x": fused_var},
|
||||
outputs={"out": fused_var},
|
||||
attrs=attrs,
|
||||
)
|
||||
|
||||
for op_idx in reversed_op_indices:
|
||||
block._remove_op(op_idx)
|
||||
|
||||
return coalesce_tensor_op_kwargs
|
||||
|
||||
|
||||
def has_same_attrs(op1, op2, attr_names):
|
||||
for attr_name in attr_names:
|
||||
if op1.attr(attr_name) != op2.attr(attr_name):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def filter_all_collective_op_indices(block):
|
||||
# NOTE: should add more collective ops
|
||||
all_collective_ops = {
|
||||
"c_broadcast",
|
||||
"broadcast",
|
||||
"all_gather",
|
||||
"all_reduce",
|
||||
}
|
||||
|
||||
match_op_indices = []
|
||||
for i, op in enumerate(block.ops):
|
||||
if op.type in all_collective_ops:
|
||||
match_op_indices.append(i)
|
||||
return match_op_indices
|
||||
|
||||
|
||||
def find_all_fuse_all_reduce_groups(block):
|
||||
collective_op_indices = filter_all_collective_op_indices(block)
|
||||
collective_ops = [block.ops[i] for i in collective_op_indices]
|
||||
|
||||
def is_valid_allreduce_op(op):
|
||||
if op.type != "c_allreduce_sum" or op.attr("use_model_parallel"):
|
||||
return False
|
||||
in_var_name = op.input("X")[0]
|
||||
out_var_name = op.output("Out")[0]
|
||||
if in_var_name != out_var_name:
|
||||
return False
|
||||
in_var = block._find_var_recursive(in_var_name)
|
||||
assert in_var is not None
|
||||
if in_var.type != core.VarDesc.VarType.DENSE_TENSOR:
|
||||
return False
|
||||
shape = in_var.shape
|
||||
if any(s <= 0 for s in shape):
|
||||
return False
|
||||
return True
|
||||
|
||||
same_attr_names = [
|
||||
"ring_id",
|
||||
"use_calc_stream",
|
||||
core.op_proto_and_checker_maker.kOpRoleAttrName(),
|
||||
core.op_proto_and_checker_maker.kOpDeviceAttrName(),
|
||||
]
|
||||
|
||||
def is_same_adjacent_op(ref_op, new_op):
|
||||
if not has_same_attrs(ref_op, new_op, same_attr_names):
|
||||
return False
|
||||
ref_op_in_var = block._find_var_recursive(ref_op.input("X")[0])
|
||||
new_op_in_var = block._find_var_recursive(new_op.input("X")[0])
|
||||
if ref_op_in_var.dtype != new_op_in_var.dtype:
|
||||
return False
|
||||
return True
|
||||
|
||||
match_seqs = find_adjacent_match_sequences(
|
||||
collective_ops, is_valid_allreduce_op, is_same_adjacent_op
|
||||
)
|
||||
new_match_seqs = []
|
||||
for i, j in match_seqs:
|
||||
new_match_seqs.append([collective_op_indices[k] for k in range(i, j)])
|
||||
return new_match_seqs
|
||||
|
||||
|
||||
def split_fuse_all_reduce_groups_by_deps(block, groups, op_deps):
|
||||
new_groups = []
|
||||
|
||||
def insert_new_group(op_indices, start_idx, end_idx):
|
||||
if end_idx - start_idx > 1:
|
||||
new_groups.append(op_indices[start_idx:end_idx])
|
||||
|
||||
for op_indices in groups:
|
||||
n = len(op_indices)
|
||||
assert n > 0
|
||||
if n == 1:
|
||||
continue
|
||||
|
||||
start_idx = 0
|
||||
k = start_idx + 1
|
||||
while k < n:
|
||||
found_group = False
|
||||
for prev_idx in range(start_idx, k):
|
||||
dep = op_deps[op_indices[prev_idx]][op_indices[k]]
|
||||
if dep == core.Node.Dep.NoDep:
|
||||
continue
|
||||
# [start_idx, k) is valid groups
|
||||
insert_new_group(op_indices, start_idx, k)
|
||||
start_idx = k
|
||||
break
|
||||
k += 1
|
||||
|
||||
insert_new_group(op_indices, start_idx, k)
|
||||
|
||||
return new_groups
|
||||
|
||||
|
||||
def insert_coalesce_tensor_ops(block, coalesce_ops_kwargs):
|
||||
if not coalesce_ops_kwargs:
|
||||
return
|
||||
|
||||
var_infos = {}
|
||||
for idx, op in enumerate(block.ops):
|
||||
for var in op.input_arg_names:
|
||||
if var not in var_infos:
|
||||
var_infos[var] = [idx, True]
|
||||
|
||||
for var in op.output_arg_names:
|
||||
if var not in var_infos:
|
||||
var_infos[var] = [idx, False]
|
||||
|
||||
n = len(block.ops)
|
||||
insert_idx_and_kwargs = []
|
||||
for group_idx, kwargs in enumerate(coalesce_ops_kwargs):
|
||||
all_vars = kwargs["inputs"]["Input"] + kwargs["outputs"]["Output"]
|
||||
min_op_idx = n
|
||||
copy_data = False
|
||||
for var in all_vars:
|
||||
if var not in var_infos:
|
||||
copy_data = True
|
||||
min_idx = 0
|
||||
break
|
||||
op_idx, is_input = var_infos[var]
|
||||
if is_input:
|
||||
copy_data = True
|
||||
min_op_idx = min(min_op_idx, op_idx)
|
||||
kwargs["attrs"]["copy_data"] = copy_data
|
||||
insert_idx_and_kwargs.append((min_op_idx, kwargs))
|
||||
|
||||
insert_idx_and_kwargs.sort(key=lambda element: element[0], reverse=True)
|
||||
for idx, kwargs in insert_idx_and_kwargs:
|
||||
block._insert_op_without_sync(idx, **kwargs)
|
||||
|
||||
|
||||
def insert_fuse_all_reduce_by_memory_size(block, groups, max_memory_size):
|
||||
op_role_key = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
||||
op_role_var_key = core.op_proto_and_checker_maker.kOpRoleVarAttrName()
|
||||
op_device_key = core.op_proto_and_checker_maker.kOpDeviceAttrName()
|
||||
coalesce_ops_kwargs = []
|
||||
for group in reversed(groups):
|
||||
first_op = block.ops[group[0]]
|
||||
ring_id = first_op.attr("ring_id")
|
||||
use_calc_stream = first_op.attr("use_calc_stream")
|
||||
use_model_parallel = first_op.attr("use_model_parallel")
|
||||
op_role = first_op.attr(op_role_key)
|
||||
op_device = first_op.attr(op_device_key)
|
||||
|
||||
attrs = {
|
||||
"ring_id": ring_id,
|
||||
"use_calc_stream": use_calc_stream,
|
||||
"use_model_parallel": use_model_parallel,
|
||||
op_role_key: op_role,
|
||||
op_device_key: op_device,
|
||||
}
|
||||
dtype = block._find_var_recursive(first_op.input("X")[0]).dtype
|
||||
sizeof = core.size_of_dtype(dtype)
|
||||
|
||||
cur_mem_size = 0
|
||||
op_role_vars = []
|
||||
recorded_op_indices = []
|
||||
in_var_names = []
|
||||
out_var_names = []
|
||||
for op_idx in reversed(group):
|
||||
op = block.ops[op_idx]
|
||||
in_var_name = op.input("X")[0]
|
||||
out_var_name = op.output("Out")[0]
|
||||
in_var = block._find_var_recursive(in_var_name)
|
||||
mem_size = int(np.prod(in_var.shape)) * sizeof
|
||||
if cur_mem_size + mem_size > max_memory_size:
|
||||
if len(recorded_op_indices) > 1:
|
||||
attrs[op_role_var_key] = op_role_vars
|
||||
coalesce_op_kwargs = insert_fuse_all_reduce_ops(
|
||||
block,
|
||||
recorded_op_indices,
|
||||
in_var_names,
|
||||
out_var_names,
|
||||
dtype,
|
||||
attrs,
|
||||
)
|
||||
coalesce_ops_kwargs.append(coalesce_op_kwargs)
|
||||
|
||||
cur_mem_size = 0
|
||||
op_role_vars = []
|
||||
recorded_op_indices = []
|
||||
in_var_names = []
|
||||
out_var_names = []
|
||||
|
||||
cur_mem_size += mem_size
|
||||
recorded_op_indices.append(op_idx)
|
||||
in_var_names.append(in_var_name)
|
||||
out_var_names.append(out_var_name)
|
||||
if op.has_attr(op_role_var_key):
|
||||
op_role_vars.extend(op.attr(op_role_var_key))
|
||||
|
||||
if len(recorded_op_indices) > 1:
|
||||
attrs[op_role_var_key] = op_role_vars
|
||||
coalesce_op_kwargs = insert_fuse_all_reduce_ops(
|
||||
block,
|
||||
recorded_op_indices,
|
||||
in_var_names,
|
||||
out_var_names,
|
||||
dtype,
|
||||
attrs,
|
||||
)
|
||||
coalesce_ops_kwargs.append(coalesce_op_kwargs)
|
||||
block._sync_with_cpp()
|
||||
insert_coalesce_tensor_ops(block, coalesce_ops_kwargs)
|
||||
|
||||
|
||||
@register_pass("fuse_all_reduce")
|
||||
class FuseAllReducePass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("max_memory_size", -1)
|
||||
|
||||
def _check_self(self):
|
||||
max_memory_size = self.get_attr("max_memory_size")
|
||||
return max_memory_size > 0
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _type(self):
|
||||
return PassType.COMM_OPT
|
||||
|
||||
# NOTE: why FuseAllReducePass can override apply_single_impl instead of
|
||||
# apply_impl? AllReduce is a collective operation, so the program of each
|
||||
# rank inside the same communication group should have the same
|
||||
# all_reduce sum operations. Therefore, FuseAllReducePass can override
|
||||
# apply_single_impl directly.
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
max_memory_size = self.get_attr("max_memory_size")
|
||||
op_deps = main_program.desc.get_op_deps()
|
||||
num_blocks = main_program.num_blocks
|
||||
for i in range(num_blocks):
|
||||
block = main_program.block(i)
|
||||
groups = find_all_fuse_all_reduce_groups(block)
|
||||
groups = split_fuse_all_reduce_groups_by_deps(
|
||||
block, groups, op_deps[i]
|
||||
)
|
||||
insert_fuse_all_reduce_by_memory_size(
|
||||
block, groups, max_memory_size
|
||||
)
|
||||
main_program._sync_with_cpp()
|
||||
+378
@@ -0,0 +1,378 @@
|
||||
# Copyright (c) 2021 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.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from paddle.framework import _apply_pass as _apply_cpp_pass
|
||||
|
||||
|
||||
class PassContext:
|
||||
def __init__(self):
|
||||
self._applied_passes = []
|
||||
self._attrs = {}
|
||||
|
||||
def set_attr(self, key, value):
|
||||
self._attrs[key] = value
|
||||
|
||||
def get_attr(self, key, default=None):
|
||||
return self._attrs.get(key, default)
|
||||
|
||||
@property
|
||||
def passes(self):
|
||||
return self._applied_passes
|
||||
|
||||
def _add_pass(self, pass_obj):
|
||||
self._applied_passes.append(pass_obj)
|
||||
|
||||
def _pop_pass(self):
|
||||
del self._applied_passes[-1]
|
||||
|
||||
|
||||
class PassType:
|
||||
UNKNOWN = 0
|
||||
COMM_OPT = 1
|
||||
CALC_OPT = 2
|
||||
PARALLEL_OPT = 3
|
||||
FUSION_OPT = 4
|
||||
|
||||
|
||||
class PassBase(ABC):
|
||||
_REGISTERED_PASSES = {}
|
||||
_COMMON_RULES = []
|
||||
|
||||
_BEFORE_WHITE_LISTS_DICT = {}
|
||||
_AFTER_WHITE_LISTS_DICT = {}
|
||||
_PASS_PROCESS_ORDER_LIST = []
|
||||
|
||||
name = None
|
||||
|
||||
@staticmethod
|
||||
def _register(pass_name, pass_class):
|
||||
assert issubclass(pass_class, PassBase)
|
||||
PassBase._REGISTERED_PASSES[pass_name] = pass_class
|
||||
|
||||
def __init__(self):
|
||||
self._attrs = {}
|
||||
|
||||
def set_attr(self, key, value):
|
||||
self._attrs[key] = value
|
||||
return self
|
||||
|
||||
def get_attr(self, key, default=None):
|
||||
return self._attrs.get(key, default)
|
||||
|
||||
@abstractmethod
|
||||
def _check_self(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _check_conflict(self, other_pass):
|
||||
pass
|
||||
|
||||
def _type(self):
|
||||
return PassType.UNKNOWN
|
||||
|
||||
def _check_conflict_including_common_rules(self, other_pass):
|
||||
return self._check_conflict(other_pass) and all(
|
||||
r(other_pass, self) for r in PassBase._COMMON_RULES
|
||||
)
|
||||
|
||||
def apply(self, main_programs, startup_programs, context=None):
|
||||
if context is None:
|
||||
context = PassContext()
|
||||
|
||||
if not self._check_self():
|
||||
return context
|
||||
|
||||
if not all(
|
||||
self._check_conflict_including_common_rules(p)
|
||||
for p in context.passes
|
||||
):
|
||||
return context
|
||||
|
||||
assert isinstance(main_programs, list)
|
||||
assert isinstance(startup_programs, list)
|
||||
assert len(main_programs) == len(startup_programs)
|
||||
self._apply_impl(main_programs, startup_programs, context)
|
||||
context._add_pass(self)
|
||||
return context
|
||||
|
||||
def _apply_impl(self, main_programs, startup_programs, context):
|
||||
for main_program, startup_program in zip(
|
||||
main_programs, startup_programs
|
||||
):
|
||||
self._apply_single_impl(main_program, startup_program, context)
|
||||
|
||||
@abstractmethod
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
pass
|
||||
|
||||
|
||||
def register_pass(name):
|
||||
def impl(cls):
|
||||
PassBase._register(name, cls)
|
||||
cls.name = name
|
||||
return cls
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def new_pass(name, pass_attrs={}):
|
||||
pass_class = PassBase._REGISTERED_PASSES.get(name)
|
||||
assert pass_class is not None, f"Pass {name} is not registered"
|
||||
pass_obj = pass_class()
|
||||
for k, v in pass_attrs.items():
|
||||
pass_obj.set_attr(k, v)
|
||||
return pass_obj
|
||||
|
||||
|
||||
class CPPPassWrapper(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def cpp_name(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def cpp_attr_types(self):
|
||||
return {}
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
_apply_cpp_pass(
|
||||
main_program,
|
||||
startup_program,
|
||||
self.cpp_name,
|
||||
self._attrs,
|
||||
self.cpp_attr_types,
|
||||
)
|
||||
|
||||
|
||||
def _fusion_opt_last_rule(pass_before, pass_after):
|
||||
if (
|
||||
pass_before._type() == PassType.FUSION_OPT
|
||||
and pass_after._type() != PassType.FUSION_OPT
|
||||
):
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
def _fusion_opt_list_rule(pass_before, pass_after):
|
||||
if (
|
||||
pass_before._type() == PassType.FUSION_OPT
|
||||
and pass_after._type() == PassType.FUSION_OPT
|
||||
):
|
||||
return _get_list_index(pass_before) < _get_list_index(pass_after)
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
def _make_rule_from_white_lists_dict(
|
||||
before_white_lists_dict, after_white_lists_dict
|
||||
):
|
||||
def collect_pass_names(white_lists_dict, result):
|
||||
for k, v in white_lists_dict.items():
|
||||
result.add(k)
|
||||
assert isinstance(v, (list, tuple))
|
||||
for pass_name in v:
|
||||
assert isinstance(pass_name, (bytes, str))
|
||||
result.add(pass_name)
|
||||
|
||||
all_pass_names = set()
|
||||
collect_pass_names(before_white_lists_dict, all_pass_names)
|
||||
collect_pass_names(after_white_lists_dict, all_pass_names)
|
||||
|
||||
compatible_pass_dict = {}
|
||||
for pass_name in all_pass_names:
|
||||
compatible_pass_dict[pass_name] = set()
|
||||
|
||||
for k, v in before_white_lists_dict.items():
|
||||
for pass_name in v:
|
||||
compatible_pass_dict[k].add(pass_name)
|
||||
|
||||
for k, v in after_white_lists_dict.items():
|
||||
for pass_name in v:
|
||||
compatible_pass_dict[pass_name].add(k)
|
||||
|
||||
def rule(pass_before, pass_after):
|
||||
all_passes_after = compatible_pass_dict.get(pass_before.name)
|
||||
if (
|
||||
all_passes_after is None
|
||||
or pass_after.name not in compatible_pass_dict
|
||||
):
|
||||
return True
|
||||
else:
|
||||
return pass_after.name in all_passes_after
|
||||
|
||||
return rule
|
||||
|
||||
|
||||
def _get_list_index(in_pass):
|
||||
assert in_pass.name in PassBase._PASS_PROCESS_ORDER_LIST, (
|
||||
f"Pass {in_pass.name} is not in _PASS_PROCESS_ORDER_LIST"
|
||||
)
|
||||
return PassBase._PASS_PROCESS_ORDER_LIST.index(in_pass.name)
|
||||
|
||||
|
||||
# The key-value pair (k, [v1, v2, ..., vn]) means the pass k can be
|
||||
# applied before any of pass [v1, v2, ..., vn] is applied
|
||||
PassBase._BEFORE_WHITE_LISTS_DICT = {
|
||||
"fuse_gradient_merge": ["fuse_all_reduce"],
|
||||
# Add more white lists here
|
||||
}
|
||||
|
||||
# The key-value pair (k, [v1, v2, ..., vn]) means the pass k can be
|
||||
# applied after any of pass [v1, v2, ..., vn] is applied
|
||||
PassBase._AFTER_WHITE_LISTS_DICT = {
|
||||
# Add more white lists here
|
||||
}
|
||||
|
||||
# The index of pass in this list represent the order in which the pass is processed.
|
||||
PassBase._PASS_PROCESS_ORDER_LIST = [
|
||||
"fuse_resunit",
|
||||
"fuse_relu_depthwise_conv",
|
||||
"fuse_bn_add_act",
|
||||
"fuse_bn_act",
|
||||
"fused_attention",
|
||||
"fused_feedforward",
|
||||
"fuse_gemm_epilogue",
|
||||
"fuse_adamw",
|
||||
"fuse_optimizer",
|
||||
]
|
||||
|
||||
PassBase._COMMON_RULES = [
|
||||
_fusion_opt_last_rule,
|
||||
_fusion_opt_list_rule,
|
||||
lambda pass_before, pass_after: type(pass_before) != type(pass_after),
|
||||
_make_rule_from_white_lists_dict(
|
||||
PassBase._BEFORE_WHITE_LISTS_DICT, PassBase._AFTER_WHITE_LISTS_DICT
|
||||
),
|
||||
# Add more common rules here
|
||||
]
|
||||
|
||||
|
||||
def _find_longest_path(edges):
|
||||
n = len(edges)
|
||||
paths = [None] * n
|
||||
dists = [None] * n
|
||||
|
||||
min_path = []
|
||||
min_dist = 0
|
||||
for i in range(n):
|
||||
paths[i] = [None] * n
|
||||
dists[i] = [None] * n
|
||||
for j in range(n):
|
||||
assert isinstance(edges[i][j], bool)
|
||||
if not edges[i][j]:
|
||||
dists[i][j] = n # inf
|
||||
paths[i][j] = []
|
||||
else:
|
||||
assert edges[i][j] is True
|
||||
dists[i][j] = -1
|
||||
paths[i][j] = [i, j]
|
||||
if dists[i][j] < min_dist:
|
||||
min_dist = -1
|
||||
min_path = paths[i][j]
|
||||
|
||||
for k in range(n):
|
||||
for i in range(n):
|
||||
for j in range(n):
|
||||
if dists[i][j] > dists[i][k] + dists[k][j]:
|
||||
dists[i][j] = dists[i][k] + dists[k][j]
|
||||
if paths[i][k]:
|
||||
assert paths[i][k][-1] == k
|
||||
else:
|
||||
continue
|
||||
if paths[k][j]:
|
||||
assert paths[k][j][0] == k
|
||||
else:
|
||||
continue
|
||||
paths[i][j] = (
|
||||
paths[i][k] + paths[k][j][1:] if paths[k][j] else []
|
||||
)
|
||||
if dists[i][j] < min_dist:
|
||||
min_dist = dists[i][j]
|
||||
min_path = paths[i][j]
|
||||
|
||||
return min_path if min_path else [0]
|
||||
|
||||
|
||||
def _solve_pass_conflict(passes, context):
|
||||
passes = [p for p in passes if p._check_self()]
|
||||
if not passes:
|
||||
return []
|
||||
|
||||
old_passes = passes
|
||||
passes = []
|
||||
for p in old_passes:
|
||||
if all(
|
||||
p._check_conflict_including_common_rules(applied_p)
|
||||
for applied_p in context.passes
|
||||
):
|
||||
passes.append(p)
|
||||
|
||||
if not passes:
|
||||
return []
|
||||
|
||||
n = len(passes)
|
||||
adjacent_matrix = []
|
||||
for _ in range(n):
|
||||
adjacent_matrix.append([None] * n)
|
||||
|
||||
for i in range(n):
|
||||
for j in range(n):
|
||||
adjacent_matrix[i][j] = passes[
|
||||
j
|
||||
]._check_conflict_including_common_rules(passes[i])
|
||||
|
||||
longest_path = _find_longest_path(adjacent_matrix)
|
||||
return [passes[idx] for idx in longest_path]
|
||||
|
||||
|
||||
class PassManager:
|
||||
def __init__(self, passes, context=None, auto_solve_conflict=True):
|
||||
if context is None:
|
||||
context = PassContext()
|
||||
self._context = context
|
||||
|
||||
if auto_solve_conflict:
|
||||
self._passes = _solve_pass_conflict(passes, context)
|
||||
else:
|
||||
self._passes = list(passes)
|
||||
|
||||
def apply(self, main_programs, startup_programs):
|
||||
context = self._context
|
||||
for p in self._passes:
|
||||
context = p.apply(main_programs, startup_programs, context)
|
||||
self._context = context
|
||||
return context
|
||||
|
||||
@property
|
||||
def context(self):
|
||||
return self._context
|
||||
|
||||
@property
|
||||
def names(self):
|
||||
return [p.name for p in self.passes]
|
||||
|
||||
@property
|
||||
def passes(self):
|
||||
return tuple(self._passes)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,54 @@
|
||||
# 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 os
|
||||
|
||||
from ..pass_base import PassContext, new_pass
|
||||
from .pipeline_1f1b import Pipeline1F1BPass # noqa: F401
|
||||
from .pipeline_eager_1f1b import PipelineEager1F1BPass # noqa: F401
|
||||
from .pipeline_fthenb import PipelineFThenBPass # noqa: F401
|
||||
from .pipeline_vpp import PipelineVirtualPipelinePass # noqa: F401
|
||||
from .pipeline_zero_bubble import (
|
||||
PipelineZeroBubblePipelinePass, # noqa: F401
|
||||
PipelineZeroBubbleVirtualPipelinePass, # noqa: F401
|
||||
)
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
def apply_pass(main_program, startup_program, pass_name, pass_attr={}):
|
||||
assert pass_name in [
|
||||
"FThenB",
|
||||
"1F1B",
|
||||
"Eager1F1B",
|
||||
"VPP",
|
||||
"ZBH1",
|
||||
"ZBVPP",
|
||||
], (
|
||||
f"pipeline scheduler only support FThenB, 1F1B, Eager1F1B, VPP and ZBH1, but receive {pass_name}"
|
||||
)
|
||||
|
||||
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([main_program], [startup_program], pass_context)
|
||||
plan = pass_context.get_attr("plan")
|
||||
return plan
|
||||
@@ -0,0 +1,341 @@
|
||||
# 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
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
from paddle.framework import (
|
||||
_current_expected_place_ as _get_device,
|
||||
)
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from ..pass_base import register_pass
|
||||
from ..pass_utils import (
|
||||
AutoParallelStreamType,
|
||||
forward_complete_op_role,
|
||||
split_program,
|
||||
)
|
||||
from .pipeline_pass_base import PipelinePassBase
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
@register_pass("pipeline_scheduler_1F1B")
|
||||
class Pipeline1F1BPass(PipelinePassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.jobs_in_stable_phase = [self.BACKWARD, self.FORWARD]
|
||||
self.jobs_in_stable_phase_in_pir = [
|
||||
self.BACKWARD,
|
||||
self.RECV_FORWARD,
|
||||
self.SEND_BACKWARD,
|
||||
self.FORWARD,
|
||||
]
|
||||
self.set_attr("enable_backward_forward_overlap", 0)
|
||||
|
||||
def _create_job_list(self):
|
||||
if self._in_pir_mode:
|
||||
return self._create_job_list_in_pir()
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"_create_job_list() only support PIR now."
|
||||
)
|
||||
|
||||
def _create_job_list_in_pir(self):
|
||||
num_micro_batches = self.get_attr("num_micro_batches")
|
||||
pp_stage = self.get_attr("pp_stage")
|
||||
pp_degree = self.get_attr("pp_degree")
|
||||
|
||||
job_list = []
|
||||
assert pp_degree <= num_micro_batches, (
|
||||
"Num of micro batches should larger than or equal to pp degree."
|
||||
)
|
||||
|
||||
micro_batch_in_warmup = pp_degree - pp_stage
|
||||
micro_batch_in_1f1b = num_micro_batches - micro_batch_in_warmup
|
||||
|
||||
forward_micro_batch_id = 0
|
||||
for i in range(micro_batch_in_warmup):
|
||||
recv_fwd_job = core.Job(self.RECV_FORWARD)
|
||||
recv_fwd_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(recv_fwd_job)
|
||||
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(forward_job)
|
||||
forward_micro_batch_id += 1
|
||||
|
||||
backward_micro_batch_id = 0
|
||||
for i in range(micro_batch_in_1f1b):
|
||||
for job_type in self.jobs_in_stable_phase_in_pir:
|
||||
job = core.Job(job_type)
|
||||
micro_batch_id = (
|
||||
forward_micro_batch_id
|
||||
if job_type.startswith(self.FORWARD)
|
||||
or job_type.startswith(self.RECV_FORWARD)
|
||||
else backward_micro_batch_id
|
||||
)
|
||||
job.set_micro_batch_id(micro_batch_id)
|
||||
job_list.append(job)
|
||||
forward_micro_batch_id += 1
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
for i in range(micro_batch_in_warmup):
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_job)
|
||||
|
||||
send_bwd_job = core.Job(self.SEND_BACKWARD)
|
||||
send_bwd_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(send_bwd_job)
|
||||
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
opt_job = core.Job(self.OPT)
|
||||
opt_job.set_micro_batch_id(0)
|
||||
job_list.append(opt_job)
|
||||
return job_list
|
||||
|
||||
def _partial_programs(self, program):
|
||||
raise NotImplementedError("pipeline_1f1b_pass() only support PIR now.")
|
||||
|
||||
def _partial_pir_programs(self, program):
|
||||
enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap")
|
||||
assert not enable_send_recv_overlap, (
|
||||
"PIR does not support 1F1B with enable_send_recv_overlap yet."
|
||||
)
|
||||
|
||||
self._overlap_send_recv(program)
|
||||
forward_complete_op_role(program)
|
||||
|
||||
job_types = [
|
||||
self.RECV_FORWARD,
|
||||
self.FORWARD,
|
||||
self.BACKWARD,
|
||||
self.SEND_BACKWARD,
|
||||
self.OPT,
|
||||
]
|
||||
|
||||
programs = {}
|
||||
for job_type in job_types:
|
||||
programs[job_type] = program.clone()
|
||||
|
||||
complete_ops = program.global_block().ops
|
||||
ops_dict = {
|
||||
key: prog.global_block().ops for key, prog in programs.items()
|
||||
}
|
||||
blocks_dict = {
|
||||
key: prog.global_block() for key, prog in programs.items()
|
||||
}
|
||||
|
||||
region = "opt"
|
||||
for op_idx in range(len(complete_ops) - 1, -1, -1):
|
||||
op = complete_ops[op_idx]
|
||||
if op.op_role != -1:
|
||||
if op.op_role == 1:
|
||||
region = "bwd"
|
||||
elif op.op_role == 0:
|
||||
region = "fwd"
|
||||
elif op.op_role == 2:
|
||||
region = "opt"
|
||||
|
||||
if region == "opt":
|
||||
self._erase_op_from_other_programs(
|
||||
op_idx, self.OPT, ops_dict, job_types
|
||||
)
|
||||
elif region == "bwd" and op.name() == "pd_op.send_v2":
|
||||
self._handle_func(
|
||||
op_idx,
|
||||
self.SEND_BACKWARD,
|
||||
job_types[4:],
|
||||
complete_ops,
|
||||
ops_dict,
|
||||
blocks_dict,
|
||||
)
|
||||
self._erase_op_from_other_programs(
|
||||
op_idx, self.SEND_BACKWARD, ops_dict, job_types
|
||||
)
|
||||
elif region == "bwd" and op.name() != "pd_op.send_v2":
|
||||
self._handle_func(
|
||||
op_idx,
|
||||
self.BACKWARD,
|
||||
job_types[3:],
|
||||
complete_ops,
|
||||
ops_dict,
|
||||
blocks_dict,
|
||||
)
|
||||
self._erase_op_from_other_programs(
|
||||
op_idx, self.BACKWARD, ops_dict, job_types
|
||||
)
|
||||
elif region == "fwd" and op.name() != "pd_op.recv_v2":
|
||||
self._handle_func(
|
||||
op_idx,
|
||||
self.FORWARD,
|
||||
job_types[2:],
|
||||
complete_ops,
|
||||
ops_dict,
|
||||
blocks_dict,
|
||||
)
|
||||
self._erase_op_from_other_programs(
|
||||
op_idx, self.FORWARD, ops_dict, job_types
|
||||
)
|
||||
elif region == "fwd" and op.name() == "pd_op.recv_v2":
|
||||
self._handle_func(
|
||||
op_idx,
|
||||
self.RECV_FORWARD,
|
||||
job_types[1:],
|
||||
complete_ops,
|
||||
ops_dict,
|
||||
blocks_dict,
|
||||
)
|
||||
self._erase_op_from_other_programs(
|
||||
op_idx, self.RECV_FORWARD, ops_dict, job_types
|
||||
)
|
||||
sub_program_list = []
|
||||
for job_type in job_types:
|
||||
sub_program_list.append(programs[job_type])
|
||||
for i in range(len(job_types)):
|
||||
logger.debug(
|
||||
f"type = {job_types[i]}, sub_programs = {sub_program_list[i]}\n"
|
||||
)
|
||||
logger.debug(
|
||||
f"jobs_in_stable_phase = {self.jobs_in_stable_phase_in_pir}"
|
||||
)
|
||||
return job_types, sub_program_list
|
||||
|
||||
def _split_program_for_overlapping(self, job_type, program, split_points):
|
||||
assert job_type in [
|
||||
self.FORWARD,
|
||||
self.BACKWARD,
|
||||
], f"job_type should be one of {[self.FORWARD, self.BACKWARD]}"
|
||||
|
||||
split_programs, __, __ = split_program(program, split_points)
|
||||
|
||||
split_job_types = []
|
||||
num_split_programs = len(split_programs)
|
||||
for idx in range(num_split_programs):
|
||||
split_job_types.append(f"{job_type}(chunk{idx})")
|
||||
|
||||
return split_job_types, split_programs
|
||||
|
||||
def is_comm_op_valid_to_overlap(self, op):
|
||||
return (
|
||||
op.type == "all_reduce"
|
||||
and op.attr("reduce_type") == paddle.distributed.ReduceOp.SUM
|
||||
and op.dist_attr.execution_stream
|
||||
== AutoParallelStreamType.CALC_STREAM.value
|
||||
)
|
||||
|
||||
def _handle_func(
|
||||
self,
|
||||
op_idx,
|
||||
cur_job_type,
|
||||
suffixed_job_types,
|
||||
complete_ops,
|
||||
ops_dict,
|
||||
blocks_dict,
|
||||
):
|
||||
for idx in range(complete_ops[op_idx].num_results()):
|
||||
if self._result_is_used(suffixed_job_types, op_idx, idx, ops_dict):
|
||||
var_name = self._get_or_create_var_name(
|
||||
ops_dict[cur_job_type], op_idx, idx, complete_ops
|
||||
)
|
||||
|
||||
for job_type in suffixed_job_types:
|
||||
if self._result_is_used([job_type], op_idx, idx, ops_dict):
|
||||
self._add_dependency_if_necessary(
|
||||
ops_dict, cur_job_type, job_type, op_idx, idx, var_name
|
||||
)
|
||||
self._add_kwarg_and_replace(
|
||||
blocks_dict[job_type],
|
||||
ops_dict[job_type],
|
||||
op_idx,
|
||||
idx,
|
||||
var_name,
|
||||
)
|
||||
|
||||
def _result_is_used(self, job_types, op_idx, rst_idx, ops_dict):
|
||||
is_used = False
|
||||
for job_type in job_types:
|
||||
is_used = (
|
||||
is_used
|
||||
or ops_dict[job_type][op_idx].result(rst_idx).use_empty()
|
||||
is False
|
||||
)
|
||||
return is_used
|
||||
|
||||
def _get_or_create_var_name(
|
||||
self, cur_sub_ops, op_idx, rst_idx, complete_ops
|
||||
):
|
||||
var_name = None
|
||||
# case1: get var_name in current sub-program
|
||||
op = cur_sub_ops[op_idx]
|
||||
if op.name() == "pd_op.data" or op.name() == "builtin.parameter":
|
||||
var_name = op.result(rst_idx).name
|
||||
else:
|
||||
# case2: get var_name from shadow_output in complete program
|
||||
result_var = complete_ops[op_idx].result(rst_idx)
|
||||
shadow_output_op = None
|
||||
for used_op in result_var.all_used_ops():
|
||||
if used_op.name() == "builtin.shadow_output":
|
||||
shadow_output_op = used_op
|
||||
if shadow_output_op is not None:
|
||||
var_name = shadow_output_op.attrs()["output_name"]
|
||||
|
||||
if var_name is None:
|
||||
# case3: create var_name in current sub-program
|
||||
paddle.pir.set_insertion_point_after(op)
|
||||
var_name = f"var_{op_idx}_{complete_ops[op_idx].name()}_{rst_idx}"
|
||||
paddle._C_ops.set_persistable_value(op.result(rst_idx), var_name)
|
||||
return var_name
|
||||
|
||||
def _add_kwarg_and_replace(self, block, ops, op_idx, rst_idx, var_name):
|
||||
ori_result = ops[op_idx].result(rst_idx)
|
||||
new_result_var = block.add_kwarg(var_name, ori_result.type())
|
||||
new_result_var.place_attr = self._get_cur_place()
|
||||
new_result_var.persistable = ori_result.persistable
|
||||
ops[op_idx].result(rst_idx).replace_all_uses_with(new_result_var)
|
||||
|
||||
def _overlap_send_recv(self, program):
|
||||
for block in program.blocks:
|
||||
for op in block.ops:
|
||||
if op.name() == "pd_op.send_v2":
|
||||
op.set_bool_attr("dynamic_shape", False)
|
||||
op.set_bool_attr("use_calc_stream", True)
|
||||
ring_id = op.attrs()["ring_id"]
|
||||
op.set_execution_stream("send_recv_stream")
|
||||
op.set_scheduling_priority(0)
|
||||
elif op.name() == "pd_op.recv_v2":
|
||||
op.set_bool_attr("dynamic_shape", False)
|
||||
op.set_bool_attr("use_calc_stream", True)
|
||||
op.set_execution_stream("send_recv_stream")
|
||||
op.set_scheduling_priority(0)
|
||||
|
||||
def _erase_op_from_other_programs(
|
||||
self, op_idx, keep_job_type, ops_dict, job_types
|
||||
):
|
||||
for job_type in job_types:
|
||||
if job_type != keep_job_type:
|
||||
ops_dict[job_type][op_idx].erase()
|
||||
|
||||
def _get_cur_place(self):
|
||||
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)
|
||||
return cur_place
|
||||
@@ -0,0 +1,73 @@
|
||||
# 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 paddle.base import core
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from ..pass_base import register_pass
|
||||
from .pipeline_pass_base import PipelinePassBase
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
@register_pass("pipeline_scheduler_Eager1F1B")
|
||||
class PipelineEager1F1BPass(PipelinePassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _create_job_list(self):
|
||||
num_micro_batches = self.get_attr("num_micro_batches")
|
||||
pp_stage = self.get_attr("pp_stage")
|
||||
pp_degree = self.get_attr("pp_degree")
|
||||
|
||||
job_list = []
|
||||
assert 2 * (pp_degree - pp_stage) - 1 <= num_micro_batches, (
|
||||
"Num of micro batches should larger than 2 * (pp_degree - pp_stage) - 1."
|
||||
)
|
||||
|
||||
micro_batch_in_warmup = 2 * (pp_degree - pp_stage) - 1
|
||||
micro_batch_in_1f1b = num_micro_batches - micro_batch_in_warmup
|
||||
|
||||
forward_micro_batch_id = 0
|
||||
for _ in range(micro_batch_in_warmup):
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(forward_job)
|
||||
forward_micro_batch_id += 1
|
||||
|
||||
backward_micro_batch_id = 0
|
||||
for _ in range(micro_batch_in_1f1b):
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_job)
|
||||
backward_micro_batch_id += 1
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(forward_job)
|
||||
forward_micro_batch_id += 1
|
||||
|
||||
for _ in range(micro_batch_in_warmup):
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_job)
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
opt_job = core.Job(self.OPT)
|
||||
job_list.append(opt_job)
|
||||
return job_list
|
||||
|
||||
def _partial_programs(self, program):
|
||||
raise NotImplementedError("Not support old IR for Eager1f1b")
|
||||
@@ -0,0 +1,66 @@
|
||||
# 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 paddle.base import core
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from ..pass_base import register_pass
|
||||
from ..pass_utils import (
|
||||
_split_program_into_forward_backward_optimize,
|
||||
)
|
||||
from .pipeline_pass_base import PipelinePassBase
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
@register_pass("pipeline_scheduler_FThenB")
|
||||
class PipelineFThenBPass(PipelinePassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _create_job_list(self):
|
||||
num_micro_batches = self.get_attr("num_micro_batches")
|
||||
|
||||
job_list = []
|
||||
|
||||
for i in range(num_micro_batches):
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(i)
|
||||
job_list.append(forward_job)
|
||||
|
||||
for i in range(num_micro_batches):
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(i)
|
||||
job_list.append(backward_job)
|
||||
|
||||
opt_job = core.Job(self.OPT)
|
||||
opt_job.set_micro_batch_id(0)
|
||||
job_list.append(opt_job)
|
||||
return job_list
|
||||
|
||||
def _partial_programs(self, program):
|
||||
raise NotImplementedError(
|
||||
"pipeline_fthenb_pass() only support PIR now."
|
||||
)
|
||||
|
||||
def _partial_pir_programs(self, program):
|
||||
# NOTE: The flag "enable_send_recv_overlap" may increase the reserved memory of GPUs.
|
||||
enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap")
|
||||
types = [self.FORWARD, self.BACKWARD, self.OPT]
|
||||
sub_program_list = _split_program_into_forward_backward_optimize(
|
||||
program, enable_send_recv_overlap
|
||||
)
|
||||
return types, sub_program_list
|
||||
@@ -0,0 +1,157 @@
|
||||
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from ..pass_base import PassBase
|
||||
from ..pass_utils import (
|
||||
set_skip_gc_vars,
|
||||
)
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
class PipelinePassBase(PassBase):
|
||||
# Pipeline stages
|
||||
RECV_FORWARD = "recv_forward"
|
||||
SEND_BACKWARD = "send_backward"
|
||||
FORWARD = "forward"
|
||||
BACKWARD = "backward"
|
||||
OPT = "optimizer"
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._in_pir_mode = paddle.base.framework.get_flags(
|
||||
"FLAGS_enable_pir_api"
|
||||
)["FLAGS_enable_pir_api"]
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _create_job_list(self):
|
||||
"""
|
||||
An interface that MUST be implemented by subclasses.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _partial_programs(self, program):
|
||||
"""
|
||||
An interface that MUST be implemented by subclasses.
|
||||
The return value MUST be two lists, one is a list of types(str), another
|
||||
is a list of sub programs.
|
||||
For example:
|
||||
return [FORWARD, BACKWARD, OPT], [fwd_prog, bwd_prog, opt_prog]
|
||||
or
|
||||
return [FORWARD], [fwd_prog]
|
||||
"""
|
||||
pass
|
||||
|
||||
def _apply_impl(self, main_programs, startup_programs, context):
|
||||
for main_program, startup_program in zip(
|
||||
main_programs, startup_programs
|
||||
):
|
||||
if self._in_pir_mode:
|
||||
self._apply_pir_single_impl(
|
||||
main_program, startup_program, context
|
||||
)
|
||||
else:
|
||||
self._apply_single_impl(main_program, startup_program, context)
|
||||
|
||||
def _partial_pir_programs(self, program):
|
||||
"""
|
||||
An interface that MUST be implemented by subclasses.
|
||||
The return value MUST be two lists, one is a list of types(str), another
|
||||
is a list of sub programs.
|
||||
For example:
|
||||
return [FORWARD, BACKWARD, OPT], [fwd_prog, bwd_prog, opt_prog]
|
||||
or
|
||||
return [FORWARD], [fwd_prog]
|
||||
"""
|
||||
pass
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, context):
|
||||
"""
|
||||
The shared process is implemented in this function and new subclass only need
|
||||
to implement two interfaces above, 'create_job_list' and 'partial_programs'.
|
||||
"""
|
||||
raise NotImplementedError("Not support for old IR")
|
||||
|
||||
def _apply_pir_single_impl(self, main_program, startup_program, context):
|
||||
"""
|
||||
The shared process is implemented in this function and new subclass only need
|
||||
to implement two interfaces above, 'create_job_list' and 'partial_programs'.
|
||||
"""
|
||||
|
||||
job_types, sub_programs = self._partial_pir_programs(main_program)
|
||||
|
||||
for i in range(len(job_types)):
|
||||
logger.debug(
|
||||
f"sub_program type: {job_types[i]}, sub_program:\n{sub_programs[i]}"
|
||||
)
|
||||
|
||||
jobs = self._create_job_list()
|
||||
type_to_program = set_skip_gc_vars(
|
||||
self.get_attr("num_micro_batches"), job_types, sub_programs, jobs
|
||||
)
|
||||
|
||||
plan = core.Plan(jobs, type_to_program)
|
||||
context.set_attr("plan", plan)
|
||||
|
||||
def _add_dependency(self, recorder_op, waiter_op, name):
|
||||
'''
|
||||
Add the extra event dependency of the two operators.
|
||||
This function mainly aims for the cross-programs in pipeline parallelism,
|
||||
especial for the 'send_v2' 'recv_v2' etc.
|
||||
'''
|
||||
if not recorder_op.has_attr("force_record_event"):
|
||||
recorder_op.set_bool_attr("force_record_event", True)
|
||||
recorder_op.set_str_attr("event_to_record", name)
|
||||
waiter_op.set_str_array_attr("events_to_wait", [name])
|
||||
|
||||
def _add_dependency_if_necessary(
|
||||
self,
|
||||
type_to_ops,
|
||||
cur_job_type,
|
||||
next_job_type,
|
||||
op_idx,
|
||||
rst_idx,
|
||||
var_name,
|
||||
):
|
||||
if not (
|
||||
("backward" in cur_job_type and "send_backward" in next_job_type)
|
||||
or ("recv_forward" in cur_job_type and "forward" in next_job_type)
|
||||
):
|
||||
return
|
||||
|
||||
first_used_idx = None
|
||||
first_used_op = None
|
||||
for used_op in (
|
||||
type_to_ops[next_job_type][op_idx].result(rst_idx).all_used_ops()
|
||||
):
|
||||
used_idx = type_to_ops[next_job_type].index(used_op)
|
||||
if first_used_idx is None or used_idx < first_used_idx:
|
||||
first_used_idx = used_idx
|
||||
first_used_op = used_op
|
||||
|
||||
if first_used_op is not None:
|
||||
self._add_dependency(
|
||||
type_to_ops[cur_job_type][op_idx], first_used_op, var_name
|
||||
)
|
||||
@@ -0,0 +1,603 @@
|
||||
# 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())
|
||||
@@ -0,0 +1,903 @@
|
||||
# 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 deque
|
||||
|
||||
from paddle.base import core
|
||||
|
||||
from ...utils.log_utils import get_logger
|
||||
from ..pass_base import register_pass
|
||||
from .pipeline_pass_base import PipelinePassBase
|
||||
|
||||
logger = get_logger(logging.INFO)
|
||||
|
||||
|
||||
class PipelineZeroBubbleBase(PipelinePassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("enable_optimizer_post_validation", 0)
|
||||
|
||||
|
||||
@register_pass("pipeline_scheduler_ZBH1")
|
||||
class PipelineZeroBubblePipelinePass(PipelineZeroBubbleBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _create_job_list(self):
|
||||
num_micro_batches = self.get_attr("num_micro_batches")
|
||||
pp_stage = self.get_attr("pp_stage")
|
||||
pp_degree = self.get_attr("pp_degree")
|
||||
|
||||
job_list = []
|
||||
assert pp_degree <= num_micro_batches, (
|
||||
"Num of micro batches should larger than or equal to pp degree."
|
||||
)
|
||||
|
||||
micro_batch_in_warmup = pp_degree - pp_stage
|
||||
micro_batch_in_zero_bubble = num_micro_batches - pp_degree
|
||||
|
||||
forward_micro_batch_id = 0
|
||||
for _ in range(micro_batch_in_warmup):
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(forward_job)
|
||||
forward_micro_batch_id += 1
|
||||
|
||||
backward_micro_batch_id = 0
|
||||
for _ in range(pp_stage):
|
||||
backward_b_job = core.Job(self.BACKWARD + "_b")
|
||||
backward_b_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_b_job)
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(forward_job)
|
||||
forward_micro_batch_id += 1
|
||||
|
||||
for _ in range(micro_batch_in_zero_bubble):
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_job)
|
||||
|
||||
forward_job = core.Job(self.FORWARD)
|
||||
forward_job.set_micro_batch_id(forward_micro_batch_id)
|
||||
job_list.append(forward_job)
|
||||
|
||||
forward_micro_batch_id += 1
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
for _ in range(micro_batch_in_warmup - 1):
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_job)
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
if pp_stage > 0:
|
||||
backward_b_job = core.Job(self.BACKWARD + "_b")
|
||||
backward_b_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_b_job)
|
||||
|
||||
backward_w_job = core.Job(self.BACKWARD + "_w")
|
||||
backward_w_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_w_job)
|
||||
else:
|
||||
backward_job = core.Job(self.BACKWARD)
|
||||
backward_job.set_micro_batch_id(backward_micro_batch_id)
|
||||
job_list.append(backward_job)
|
||||
backward_micro_batch_id += 1
|
||||
|
||||
for i in range(pp_stage):
|
||||
backward_w_job = core.Job(self.BACKWARD + "_w")
|
||||
backward_w_job.set_micro_batch_id(i)
|
||||
job_list.append(backward_w_job)
|
||||
|
||||
opt_job = core.Job(self.OPT)
|
||||
opt_job.set_micro_batch_id(0)
|
||||
job_list.append(opt_job)
|
||||
return job_list
|
||||
|
||||
def _partial_programs(self, program):
|
||||
raise NotImplementedError("Not support old IR for ZeroBubble")
|
||||
|
||||
|
||||
@register_pass("pipeline_scheduler_ZBVPP")
|
||||
class PipelineZeroBubbleVirtualPipelinePass(PipelineZeroBubblePipelinePass):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.set_attr("enable_optimizer_post_validation", 0)
|
||||
self.set_attr("program_runtimes", [61, 72, 71, 34, 3])
|
||||
self.set_attr("memory_limit_times", -1)
|
||||
|
||||
self.program_mem_usages = []
|
||||
self.program_max_mem_usages = []
|
||||
self.base_memory = []
|
||||
self.program_runtime = {}
|
||||
|
||||
def _create_job_list(self):
|
||||
pp_degree = self.get_attr("pp_degree")
|
||||
num_micro_batches = self.get_attr("num_micro_batches")
|
||||
num_model_chunks = self.get_attr("vpp_degree")
|
||||
|
||||
assert num_micro_batches % pp_degree == 0
|
||||
|
||||
# TODO(luchang): Fix the gradient explosion issue when num_model_chunks(accumulate steps) > pp_degree
|
||||
assert num_micro_batches <= pp_degree, (
|
||||
"zbvpp now only supports accumulate steps <= pp degree. It will cause gradient exploitation when accumulate steps > pp degree."
|
||||
)
|
||||
|
||||
program_runtimes = self.get_attr("program_runtimes")
|
||||
|
||||
self.program_runtime = {
|
||||
"forward": program_runtimes[0],
|
||||
"backward_b": program_runtimes[1],
|
||||
"backward_w": program_runtimes[2],
|
||||
"loss": program_runtimes[3],
|
||||
"communication": program_runtimes[4],
|
||||
}
|
||||
|
||||
v_scheduler = VScheduleCreator(
|
||||
pp_degree,
|
||||
num_micro_batches,
|
||||
num_model_chunks,
|
||||
self.program_mem_usages,
|
||||
self.program_max_mem_usages,
|
||||
self.base_memory,
|
||||
self.program_runtime,
|
||||
self._get_max_memory(),
|
||||
)
|
||||
|
||||
schedule, end_time = None, None
|
||||
for fill_w_before_b in [True, False]:
|
||||
for fill_w_before_f in [True, False]:
|
||||
for fill_loss_stage in [True, False]:
|
||||
if schedule is None:
|
||||
schedule, end_time, _ = v_scheduler.create_v_schedule(
|
||||
fill_w_before_b=fill_w_before_b,
|
||||
fill_w_before_f=fill_w_before_f,
|
||||
fill_loss_stage=fill_loss_stage,
|
||||
)
|
||||
else:
|
||||
(
|
||||
new_schedule,
|
||||
new_end_time,
|
||||
_,
|
||||
) = v_scheduler.create_v_schedule(
|
||||
fill_w_before_b=fill_w_before_b,
|
||||
fill_w_before_f=fill_w_before_f,
|
||||
fill_loss_stage=fill_loss_stage,
|
||||
)
|
||||
if max(new_end_time) < max(end_time):
|
||||
schedule, end_time = new_schedule, new_end_time
|
||||
|
||||
stage_schedule = schedule[self.get_attr("pp_stage")]
|
||||
job_list = []
|
||||
|
||||
for job_info in stage_schedule:
|
||||
job = core.Job(f"{job_info['type']}{job_info['chunk']}")
|
||||
job.set_micro_batch_id(job_info["micro_batch"])
|
||||
job_list.append(job)
|
||||
|
||||
opt_job = core.Job(self.OPT)
|
||||
opt_job.set_micro_batch_id(0)
|
||||
job_list.append(opt_job)
|
||||
|
||||
return job_list
|
||||
|
||||
def _partial_programs(self, program):
|
||||
raise NotImplementedError("Not support old IR for ZeroBubbleVPP")
|
||||
|
||||
|
||||
class VScheduleCreator:
|
||||
def __init__(
|
||||
self,
|
||||
num_stage,
|
||||
num_micro_batch,
|
||||
num_model_chunks,
|
||||
program_mem_usages,
|
||||
program_max_mem_usages,
|
||||
base_memory,
|
||||
program_runtime,
|
||||
max_memory=None,
|
||||
):
|
||||
self.num_stage = num_stage
|
||||
self.num_micro_batch = num_micro_batch
|
||||
self.num_model_chunks = num_model_chunks
|
||||
self.num_nodes = num_model_chunks * num_stage * num_micro_batch * 3
|
||||
self.program_mem_usages = program_mem_usages
|
||||
self.program_max_mem_usages = program_max_mem_usages
|
||||
self.job_types = ["forward", "backward_w", "backward_b"]
|
||||
self.program_runtime = program_runtime
|
||||
self.base_memory = base_memory
|
||||
self.max_memory = max_memory
|
||||
if max_memory is None:
|
||||
self.max_memory = float("inf")
|
||||
self.calculate_loss_stage = (
|
||||
0 if num_model_chunks % 2 == 0 else num_stage - 1
|
||||
)
|
||||
|
||||
def init_schedule(self):
|
||||
job_counter = {}
|
||||
for job_type in self.job_types:
|
||||
for chunk_id in range(self.num_model_chunks):
|
||||
job_counter[f"{job_type}{chunk_id}"] = 0
|
||||
|
||||
self._job_counters = [job_counter.copy() for _ in range(self.num_stage)]
|
||||
self._job_end_times = [-1] * self.num_nodes
|
||||
self._stage_current_time = [0] * self.num_stage
|
||||
self._stage_mem_usage = self.base_memory.copy()
|
||||
self._pending_w = [deque() for _ in range(self.num_stage)]
|
||||
self._stage_job_schedule = [[] for _ in range(self.num_stage)]
|
||||
self._stage_bubbles = [0] * self.num_stage
|
||||
|
||||
def create_v_schedule(
|
||||
self,
|
||||
fill_w_before_f=True,
|
||||
fill_w_before_b=True,
|
||||
fill_loss_stage=True,
|
||||
approved_bubbles=None,
|
||||
):
|
||||
self.init_schedule()
|
||||
if approved_bubbles is None:
|
||||
approved_bubbles = [-1] * self.num_stage
|
||||
max_approved_bubble = max(approved_bubbles)
|
||||
|
||||
self._insert_forward_jobs_before_forward1()
|
||||
self._insert_forward_jobs_before_backward_b()
|
||||
self._insert_jobs_after_backward_start(
|
||||
fill_w_before_f, fill_w_before_b, fill_loss_stage, approved_bubbles
|
||||
)
|
||||
|
||||
schedule = self._stage_job_schedule.copy()
|
||||
end_time = self._job_end_times.copy()
|
||||
max_bubble = self._get_max_stage_bubble()
|
||||
|
||||
if max_approved_bubble < 0 or max_bubble < max_approved_bubble:
|
||||
new_schedule, new_end_time, new_max_bubble = self.create_v_schedule(
|
||||
fill_w_before_f,
|
||||
fill_w_before_b,
|
||||
fill_loss_stage,
|
||||
self._stage_bubbles,
|
||||
)
|
||||
|
||||
if max(new_end_time) < max(end_time):
|
||||
return new_schedule, new_end_time, new_max_bubble
|
||||
|
||||
return schedule, end_time, max_bubble
|
||||
|
||||
def _insert_forward_jobs_before_forward1(self):
|
||||
# Step1: Insert forward jobs with chunk_id=0 into the schedule
|
||||
for i in range(self.num_stage):
|
||||
self._put_job_into_schedule("forward", chunk_id=0, stage_id=i)
|
||||
|
||||
# Step2: Insert forward jobs with chunk_id=1 into the schedule
|
||||
self._fill_forward_before_one_job(
|
||||
"forward", 1, [0], forward_insert_order="down"
|
||||
)
|
||||
|
||||
def _insert_forward_jobs_before_backward_b(self):
|
||||
chunk_to_insert_order = {
|
||||
0: "down",
|
||||
1: "up",
|
||||
}
|
||||
|
||||
# Insert the rest chunk_id forward jobs
|
||||
for chunk_id in range(2, self.num_model_chunks):
|
||||
fill_chunk_ids = list(range(0, chunk_id))
|
||||
forward_insert_order = chunk_to_insert_order[fill_chunk_ids[-1] % 2]
|
||||
self._fill_forward_before_one_job(
|
||||
"forward", chunk_id, fill_chunk_ids, forward_insert_order
|
||||
)
|
||||
|
||||
# Insert forward jobs to fill the bubble before backward_b0
|
||||
fill_chunk_ids = list(range(0, self.num_model_chunks))
|
||||
forward_insert_order = chunk_to_insert_order[fill_chunk_ids[-1] % 2]
|
||||
|
||||
self._fill_forward_before_one_job(
|
||||
"backward_b",
|
||||
0,
|
||||
list(range(0, self.num_model_chunks)),
|
||||
forward_insert_order,
|
||||
insert_end_point_job=False,
|
||||
)
|
||||
|
||||
def _fill_forward_before_one_job(
|
||||
self,
|
||||
end_point_job_type,
|
||||
end_point_chunk_id,
|
||||
fill_chunk_ids,
|
||||
forward_insert_order,
|
||||
insert_end_point_job=True,
|
||||
):
|
||||
stage_order = list(range(self.num_stage))
|
||||
if forward_insert_order == "down":
|
||||
stage_order.reverse()
|
||||
|
||||
stage_last_job = self._stage_job_schedule[stage_order[0]][-1]
|
||||
end_point_job_start_time = self._job_end_times[
|
||||
self._get_job_id(
|
||||
stage_last_job["type"],
|
||||
stage_last_job["chunk"],
|
||||
stage_order[0],
|
||||
stage_last_job["micro_batch"],
|
||||
)
|
||||
]
|
||||
|
||||
if insert_end_point_job:
|
||||
self._put_job_into_schedule(
|
||||
end_point_job_type, end_point_chunk_id, stage_order[0]
|
||||
)
|
||||
|
||||
for stage_id in stage_order[1:]:
|
||||
stage_last_job = self._stage_job_schedule[stage_id][-1]
|
||||
start_fill_chunk = (stage_last_job["chunk"] + 1) % len(
|
||||
fill_chunk_ids
|
||||
)
|
||||
|
||||
end_point_job_start_time = (
|
||||
end_point_job_start_time
|
||||
+ self.program_runtime["communication"]
|
||||
+ self._get_program_runtime(
|
||||
end_point_job_type, stage_id, end_point_chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
self._fill_bubble_with_forward(
|
||||
stage_id,
|
||||
fill_chunk_ids,
|
||||
start_fill_chunk,
|
||||
end_point_job_start_time,
|
||||
insert_order=forward_insert_order,
|
||||
)
|
||||
|
||||
if insert_end_point_job:
|
||||
self._put_job_into_schedule(
|
||||
end_point_job_type, end_point_chunk_id, stage_id
|
||||
)
|
||||
|
||||
def _insert_jobs_after_backward_start(
|
||||
self,
|
||||
fill_w_before_f,
|
||||
fill_w_before_b,
|
||||
fill_loss_stage,
|
||||
approved_bubbles,
|
||||
):
|
||||
backward_b_job_number = self.num_model_chunks * self.num_micro_batch
|
||||
|
||||
first_backward_b_stage = (
|
||||
0 if self.num_model_chunks % 2 == 0 else self.num_stage - 1
|
||||
)
|
||||
while (
|
||||
self._get_stage_backward_b_number(first_backward_b_stage)
|
||||
< backward_b_job_number
|
||||
):
|
||||
# Step1: Check memory usage, if not enough, put pending backward_w job into schedule
|
||||
for stage_id in range(self.num_stage):
|
||||
while not self._memory_check("backward_b", 0, stage_id):
|
||||
if len(self._pending_w[stage_id]) == 0:
|
||||
raise ValueError(
|
||||
f"No pending backward_w job and backward_b0 job exceeds the memory limit at stage {stage_id}."
|
||||
)
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
# Step2: Insert backward_b job for each stage
|
||||
# b_ranks = [[] for _ in range(self.num_model_chunks)]
|
||||
backward_insert_order = range(self.num_stage)
|
||||
if self.num_model_chunks % 2:
|
||||
backward_insert_order = range(self.num_stage - 1, -1, -1)
|
||||
|
||||
for stage_id in backward_insert_order:
|
||||
for chunk_id in range(0, self.num_model_chunks):
|
||||
if self._can_schedule_b_task(stage_id, chunk_id):
|
||||
dependency_job_end_time = (
|
||||
self._get_dependency_job_end_time(
|
||||
"backward_b",
|
||||
chunk_id,
|
||||
stage_id,
|
||||
self._job_counters[stage_id][
|
||||
f"backward_b{chunk_id}"
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
while len(
|
||||
self._pending_w[stage_id]
|
||||
) and dependency_job_end_time + self.program_runtime[
|
||||
"communication"
|
||||
] >= self._stage_current_time[
|
||||
stage_id
|
||||
] + self._get_program_runtime(
|
||||
"backward_w",
|
||||
stage_id,
|
||||
self._pending_w[stage_id][0][0],
|
||||
):
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
if (
|
||||
stage_id == self.calculate_loss_stage
|
||||
and fill_loss_stage
|
||||
):
|
||||
while (
|
||||
len(self._pending_w[stage_id])
|
||||
and dependency_job_end_time
|
||||
+ self.program_runtime["communication"]
|
||||
>= self._stage_current_time[stage_id]
|
||||
+ self._get_program_runtime(
|
||||
"backward_w",
|
||||
stage_id,
|
||||
self._pending_w[stage_id][0][0],
|
||||
)
|
||||
* 0.2
|
||||
):
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
max_stage_bubble = self._get_max_stage_bubble(
|
||||
stage_id, approved_bubbles
|
||||
)
|
||||
stage_bubble = self._stage_bubbles[stage_id]
|
||||
|
||||
if (
|
||||
len(self._pending_w[stage_id])
|
||||
and dependency_job_end_time
|
||||
+ self.program_runtime["communication"]
|
||||
- self._stage_current_time[stage_id]
|
||||
> max_stage_bubble - stage_bubble
|
||||
):
|
||||
if fill_w_before_b:
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
self._put_job_into_schedule(
|
||||
"backward_b", chunk_id, stage_id
|
||||
)
|
||||
break
|
||||
|
||||
# Step3: Insert forward jobs after backward_b
|
||||
forward_insert_order = range(self.num_stage)
|
||||
if self.num_model_chunks % 2:
|
||||
forward_insert_order = range(self.num_stage - 1, -1, -1)
|
||||
|
||||
for stage_id in forward_insert_order:
|
||||
for chunk_id in range(self.num_model_chunks - 1, -1, -1):
|
||||
if self._can_schedule_f_task(stage_id, chunk_id):
|
||||
while (
|
||||
self._stage_mem_usage[stage_id]
|
||||
+ self.program_max_mem_usages[stage_id][
|
||||
f"forward{chunk_id}"
|
||||
]
|
||||
> self.max_memory
|
||||
):
|
||||
if len(self._pending_w[stage_id]) == 0:
|
||||
raise ValueError(
|
||||
"No pending backward_w job and forward job exceeds the memory limit."
|
||||
)
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
dependency_job_end_time = (
|
||||
self._get_dependency_job_end_time(
|
||||
"forward",
|
||||
chunk_id,
|
||||
stage_id,
|
||||
self._job_counters[stage_id][
|
||||
f"forward{chunk_id}"
|
||||
],
|
||||
)
|
||||
)
|
||||
while len(
|
||||
self._pending_w[stage_id]
|
||||
) and dependency_job_end_time + self.program_runtime[
|
||||
"communication"
|
||||
] >= self._stage_current_time[
|
||||
stage_id
|
||||
] + self._get_program_runtime(
|
||||
"backward_w",
|
||||
stage_id,
|
||||
self._pending_w[stage_id][0][0],
|
||||
):
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
max_stage_bubble = self._get_max_stage_bubble(
|
||||
stage_id, approved_bubbles
|
||||
)
|
||||
stage_bubble = self._stage_bubbles[stage_id]
|
||||
if (
|
||||
len(self._pending_w[stage_id])
|
||||
and dependency_job_end_time
|
||||
+ self.program_runtime["communication"]
|
||||
- self._stage_current_time[stage_id]
|
||||
> max_stage_bubble - stage_bubble
|
||||
):
|
||||
if fill_w_before_f:
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
self._put_job_into_schedule(
|
||||
"forward", chunk_id, stage_id
|
||||
)
|
||||
break
|
||||
|
||||
for stage_id in range(self.num_stage):
|
||||
while len(self._pending_w[stage_id]):
|
||||
self._put_w_job_into_schedule(stage_id)
|
||||
|
||||
def _can_schedule_f_task(self, stage_id, chunk_id):
|
||||
return self._can_schedule_task("forward", chunk_id, stage_id)
|
||||
|
||||
def _can_schedule_b_task(self, stage_id, chunk_id):
|
||||
return self._can_schedule_task("backward_b", chunk_id, stage_id)
|
||||
|
||||
def _can_schedule_task(self, job_type, chunk_id, stage_id):
|
||||
if job_type == "forward":
|
||||
current_key = f"forward{chunk_id}"
|
||||
prev_key = f"forward{chunk_id - 1}"
|
||||
elif job_type == "backward_b":
|
||||
current_key = f"backward_b{chunk_id}"
|
||||
prev_chunk_id = chunk_id + 1
|
||||
if prev_chunk_id >= self.num_model_chunks:
|
||||
prev_key = f"forward{self.num_model_chunks - 1}"
|
||||
else:
|
||||
prev_key = f"backward_b{chunk_id + 1}"
|
||||
|
||||
micro_batch_id = self._job_counters[stage_id][current_key]
|
||||
if micro_batch_id >= self.num_micro_batch:
|
||||
return False
|
||||
|
||||
if (job_type == "forward" and chunk_id > 0) or (
|
||||
job_type == "backward_b"
|
||||
):
|
||||
prev_chunk_count = self._job_counters[stage_id][prev_key]
|
||||
current_chunk_count = self._job_counters[stage_id][current_key]
|
||||
if prev_chunk_count <= current_chunk_count:
|
||||
return False
|
||||
|
||||
prev_stage_job_end_time = self._get_dependency_job_end_time(
|
||||
job_type, chunk_id, stage_id, micro_batch_id
|
||||
)
|
||||
|
||||
if prev_stage_job_end_time < 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _get_stage_micro_batch_id(self, stage_id, job_type):
|
||||
micro_batch_id = 0
|
||||
for chunk_id in range(self.num_model_chunks):
|
||||
micro_batch_id += self._job_counters[stage_id][
|
||||
f"{job_type}{chunk_id}"
|
||||
]
|
||||
return micro_batch_id
|
||||
|
||||
def _fill_bubble_with_forward(
|
||||
self,
|
||||
stage_id,
|
||||
chunk_ids,
|
||||
start_chunk_id,
|
||||
next_job_start_time,
|
||||
insert_order="down",
|
||||
):
|
||||
chunk_id = start_chunk_id
|
||||
less_forward_number = 0
|
||||
stage_order = range(0, stage_id + 1)
|
||||
if insert_order == "up":
|
||||
less_forward_number = self._get_stage_forward_number(0, chunk_ids)
|
||||
stage_order = range(self.num_stage - 1, stage_id - 1, -1)
|
||||
|
||||
while self._check_before_insert(
|
||||
chunk_ids, stage_order, next_job_start_time, less_forward_number
|
||||
):
|
||||
available_memory = self.max_memory - self._stage_mem_usage[stage_id]
|
||||
# After insert forward job, we need to check whether we can insert backward_b job
|
||||
available_memory -= self.program_max_mem_usages[stage_id][
|
||||
f"backward_b{chunk_id}"
|
||||
]
|
||||
|
||||
# Check whether we can insert all chunk_id forward jobs
|
||||
for i in range(1, self.num_model_chunks):
|
||||
if self._job_counters[stage_id][f"forward{i}"] == 0:
|
||||
available_memory -= self.program_max_mem_usages[stage_id][
|
||||
f"forward{i}"
|
||||
]
|
||||
if (
|
||||
available_memory
|
||||
< self.program_max_mem_usages[stage_id][f"forward{chunk_id}"]
|
||||
):
|
||||
break
|
||||
|
||||
for i in stage_order:
|
||||
if self._can_schedule_f_task(i, chunk_id):
|
||||
stage_forward_number = self._get_stage_forward_number(
|
||||
i, chunk_ids
|
||||
)
|
||||
if stage_forward_number >= less_forward_number:
|
||||
if not self._time_check(
|
||||
"forward", chunk_id, i, next_job_start_time
|
||||
):
|
||||
continue
|
||||
self._put_job_into_schedule("forward", chunk_id, i)
|
||||
|
||||
chunk_id = (chunk_id + 1) % len(chunk_ids)
|
||||
|
||||
def _check_before_insert(
|
||||
self, chunk_ids, stage_order, next_job_start_time, less_forward_number
|
||||
):
|
||||
stage_id = stage_order[-1]
|
||||
if (
|
||||
self._get_stage_forward_number(stage_id, chunk_ids)
|
||||
< less_forward_number
|
||||
):
|
||||
return True
|
||||
|
||||
job_numbers = []
|
||||
for chunk_id in chunk_ids:
|
||||
job_numbers.append(
|
||||
self._job_counters[stage_id][f"forward{chunk_id}"]
|
||||
)
|
||||
|
||||
micro_batch_id_check = False
|
||||
for number in job_numbers:
|
||||
if number < self.num_micro_batch:
|
||||
micro_batch_id_check = True
|
||||
break
|
||||
|
||||
can_insert = False
|
||||
for chunk_id in chunk_ids:
|
||||
for i in stage_order:
|
||||
if self._can_schedule_f_task(i, chunk_id):
|
||||
stage_forward_number = self._get_stage_forward_number(
|
||||
i, chunk_ids
|
||||
)
|
||||
if stage_forward_number >= less_forward_number:
|
||||
if not self._time_check(
|
||||
"forward", chunk_id, i, next_job_start_time
|
||||
):
|
||||
continue
|
||||
can_insert = True
|
||||
break
|
||||
|
||||
return (
|
||||
self._memory_check("forward", chunk_id, stage_id)
|
||||
and micro_batch_id_check
|
||||
and can_insert
|
||||
)
|
||||
|
||||
def _memory_check(self, job_type, chunk_id, stage_id):
|
||||
if (
|
||||
self._stage_mem_usage[stage_id]
|
||||
+ self.program_max_mem_usages[stage_id][f"{job_type}{chunk_id}"]
|
||||
> self.max_memory
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _time_check(self, job_type, chunk_id, stage_id, next_job_start_time):
|
||||
dependency_job_end_time = self._get_dependency_job_end_time(
|
||||
job_type,
|
||||
chunk_id,
|
||||
stage_id,
|
||||
self._job_counters[stage_id][f"{job_type}{chunk_id}"],
|
||||
)
|
||||
job_end_time = max(
|
||||
self._stage_current_time[stage_id],
|
||||
dependency_job_end_time + self.program_runtime["communication"],
|
||||
) + self._get_program_runtime(job_type, stage_id, chunk_id)
|
||||
|
||||
if job_end_time > next_job_start_time:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _micro_batch_id_check(self, job_type, chunk_id, stage_id):
|
||||
if (
|
||||
self._job_counters[stage_id][f"{job_type}{chunk_id}"]
|
||||
>= self.num_micro_batch
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _put_job_into_schedule(
|
||||
self,
|
||||
job_type,
|
||||
chunk_id,
|
||||
stage_id,
|
||||
):
|
||||
program_runtime = self._get_program_runtime(
|
||||
job_type, stage_id, chunk_id
|
||||
)
|
||||
task_end_time = self._stage_current_time[stage_id] + program_runtime
|
||||
|
||||
micro_batch_id = self._job_counters[stage_id][f"{job_type}{chunk_id}"]
|
||||
if micro_batch_id >= self.num_micro_batch:
|
||||
raise ValueError(
|
||||
f"Job {job_type}{chunk_id} exceeds the limit of micro batches."
|
||||
)
|
||||
|
||||
if (
|
||||
self._stage_mem_usage[stage_id]
|
||||
+ self.program_max_mem_usages[stage_id][f"{job_type}{chunk_id}"]
|
||||
> self.max_memory
|
||||
):
|
||||
raise ValueError(
|
||||
f"Job {job_type}{chunk_id} exceeds the memory limit at stage {stage_id}."
|
||||
)
|
||||
|
||||
self._check_job_chunk_order(
|
||||
job_type, chunk_id, stage_id, micro_batch_id
|
||||
)
|
||||
|
||||
if job_type in ["forward", "backward_b"]:
|
||||
dependency_job_end_time = self._get_dependency_job_end_time(
|
||||
job_type, chunk_id, stage_id, micro_batch_id
|
||||
)
|
||||
if dependency_job_end_time < 0:
|
||||
prev_stage_id = self._get_prev_stage_id(
|
||||
job_type, chunk_id, stage_id
|
||||
)
|
||||
raise ValueError(
|
||||
f"Job {job_type}{chunk_id}_{micro_batch_id} at stage {stage_id} depends on unscheduled job {job_type}{chunk_id}_{micro_batch_id} at stage {prev_stage_id}."
|
||||
)
|
||||
task_end_time = max(
|
||||
task_end_time,
|
||||
dependency_job_end_time
|
||||
+ self.program_runtime["communication"]
|
||||
+ program_runtime,
|
||||
)
|
||||
|
||||
job_id = self._get_job_id(job_type, chunk_id, stage_id, micro_batch_id)
|
||||
if self._job_counters[stage_id]["forward0"] > 0:
|
||||
self._stage_bubbles[stage_id] += (
|
||||
task_end_time
|
||||
- self._stage_current_time[stage_id]
|
||||
- program_runtime
|
||||
)
|
||||
|
||||
self._job_end_times[job_id] = task_end_time
|
||||
self._stage_current_time[stage_id] = task_end_time
|
||||
self._stage_mem_usage[stage_id] += self.program_mem_usages[stage_id][
|
||||
f"{job_type}{chunk_id}"
|
||||
]
|
||||
|
||||
job_info = {
|
||||
"type": job_type,
|
||||
"chunk": chunk_id,
|
||||
"micro_batch": micro_batch_id,
|
||||
}
|
||||
self._stage_job_schedule[stage_id].append(job_info)
|
||||
if job_type == "backward_b":
|
||||
self._pending_w[stage_id].append((chunk_id, micro_batch_id))
|
||||
self._job_counters[stage_id][f"{job_type}{chunk_id}"] += 1
|
||||
|
||||
def _put_w_job_into_schedule(self, stage_id):
|
||||
if not len(self._pending_w[stage_id]):
|
||||
raise ValueError("No pending backward_w job.")
|
||||
|
||||
chunk_id, _ = self._pending_w[stage_id].popleft()
|
||||
self._put_job_into_schedule("backward_w", chunk_id, stage_id)
|
||||
|
||||
def _check_job_chunk_order(
|
||||
self, job_type, chunk_id, stage_id, micro_batch_id
|
||||
):
|
||||
if job_type == "forward":
|
||||
if chunk_id > 0:
|
||||
prev_job_end_time = self._job_end_times[
|
||||
self._get_job_id(
|
||||
"forward", chunk_id - 1, stage_id, micro_batch_id
|
||||
)
|
||||
]
|
||||
if prev_job_end_time < 0:
|
||||
raise ValueError(
|
||||
f"Job {job_type}{chunk_id}_{micro_batch_id} depends on unfinished {job_type}{chunk_id - 1}_{micro_batch_id} job."
|
||||
)
|
||||
elif job_type == "backward_b":
|
||||
if chunk_id < self.num_model_chunks - 1:
|
||||
prev_job_end_time = self._job_end_times[
|
||||
self._get_job_id(
|
||||
job_type, chunk_id + 1, stage_id, micro_batch_id
|
||||
)
|
||||
]
|
||||
if prev_job_end_time < 0:
|
||||
raise ValueError(
|
||||
f"Job {job_type}{chunk_id}_{micro_batch_id} depends on unfinished {job_type}{chunk_id + 1}_{micro_batch_id} job."
|
||||
)
|
||||
elif job_type == "backward_w":
|
||||
prev_job_id = self._get_job_id(
|
||||
"backward_b", chunk_id, stage_id, micro_batch_id
|
||||
)
|
||||
if self._job_end_times[prev_job_id] < 0:
|
||||
raise ValueError(
|
||||
f"Job {job_type}{chunk_id}_{micro_batch_id} at stage {stage_id} depends on unfinished backward_b{chunk_id}_{micro_batch_id} job."
|
||||
)
|
||||
|
||||
def _get_dependency_job_end_time(
|
||||
self, job_type, chunk_id, stage_id, micro_batch_id
|
||||
):
|
||||
prev_stage_id = self._get_prev_stage_id(job_type, chunk_id, stage_id)
|
||||
if prev_stage_id < 0 or prev_stage_id >= self.num_stage:
|
||||
return 0
|
||||
|
||||
prev_stage_job_id = self._get_job_id(
|
||||
job_type, chunk_id, prev_stage_id, micro_batch_id
|
||||
)
|
||||
|
||||
prev_job_end_time = self._job_end_times[prev_stage_job_id]
|
||||
return prev_job_end_time
|
||||
|
||||
def _get_prev_stage_id(self, job_type, chunk_id, stage_id):
|
||||
if job_type == "forward":
|
||||
if chunk_id % 2:
|
||||
return stage_id + 1
|
||||
else:
|
||||
return stage_id - 1
|
||||
elif job_type in ["backward_b", "backward_w"]:
|
||||
if chunk_id % 2:
|
||||
return stage_id - 1
|
||||
else:
|
||||
return stage_id + 1
|
||||
|
||||
def _get_max_stage_bubble(self, stage_id=-1, approved_bubbles=None):
|
||||
max_stage_bubble = max(self._stage_bubbles)
|
||||
if stage_id >= 0:
|
||||
max_approved_bubble = max(approved_bubbles)
|
||||
max_stage_bubble = max(
|
||||
max_stage_bubble,
|
||||
max_approved_bubble - approved_bubbles[stage_id],
|
||||
)
|
||||
return max_stage_bubble
|
||||
|
||||
def _get_job_id(self, job_type, chunk_id, stage_id, job_micro_id):
|
||||
return (
|
||||
self.job_types.index(job_type)
|
||||
* self.num_model_chunks
|
||||
* self.num_stage
|
||||
* self.num_micro_batch
|
||||
+ chunk_id * self.num_stage * self.num_micro_batch
|
||||
+ stage_id * self.num_micro_batch
|
||||
+ job_micro_id
|
||||
)
|
||||
|
||||
def _get_bubble_rate(self):
|
||||
max_bubble = self._get_max_stage_bubble()
|
||||
fbw_cost = (
|
||||
self.program_runtime["forward"]
|
||||
+ self.program_runtime["backward_w"]
|
||||
+ self.program_runtime["communication"]
|
||||
)
|
||||
expected_time = fbw_cost * self.num_micro_batch * self.num_model_chunks
|
||||
bubble_rate = max_bubble / expected_time
|
||||
return bubble_rate
|
||||
|
||||
def _get_stage_forward_number(self, stage, chunk_ids=None):
|
||||
job_number = 0
|
||||
if chunk_ids is None:
|
||||
chunk_ids = range(self.num_model_chunks)
|
||||
for chunk_id in chunk_ids:
|
||||
job_number += self._job_counters[stage][f"forward{chunk_id}"]
|
||||
return job_number
|
||||
|
||||
def _get_stage_backward_b_number(self, stage, chunk_ids=None):
|
||||
job_number = 0
|
||||
if chunk_ids is None:
|
||||
chunk_ids = range(self.num_model_chunks)
|
||||
for chunk_id in chunk_ids:
|
||||
job_number += self._job_counters[stage][f"backward_b{chunk_id}"]
|
||||
return job_number
|
||||
|
||||
def _get_program_runtime(self, job_type, stage_id, chunk_id):
|
||||
program_runtime = self.program_runtime[job_type]
|
||||
|
||||
if job_type == "communication":
|
||||
return program_runtime
|
||||
|
||||
if (
|
||||
stage_id == self.calculate_loss_stage
|
||||
and chunk_id == self.num_model_chunks - 1
|
||||
):
|
||||
program_runtime += self.program_runtime["loss"]
|
||||
return program_runtime
|
||||
+274
@@ -0,0 +1,274 @@
|
||||
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
|
||||
import paddle
|
||||
from paddle.optimizer.lr import (
|
||||
ExponentialDecay,
|
||||
InverseTimeDecay,
|
||||
LRScheduler,
|
||||
NaturalExpDecay,
|
||||
NoamDecay,
|
||||
exponential_decay,
|
||||
inverse_time_decay,
|
||||
noam_decay,
|
||||
)
|
||||
|
||||
from ..ps.utils.public import (
|
||||
get_optimize_ops,
|
||||
get_ps_endpoint,
|
||||
get_role_id,
|
||||
get_trainers,
|
||||
)
|
||||
from .pass_base import PassBase, register_pass
|
||||
|
||||
|
||||
@register_pass("add_lr_decay_table_pass")
|
||||
class AddLrDecayTablePass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _add_tensor_table(
|
||||
self,
|
||||
attrs,
|
||||
feed_var_name,
|
||||
fetch_var_name="",
|
||||
startup_program=None,
|
||||
main_program=None,
|
||||
tensor_table_class="",
|
||||
):
|
||||
tensor_table_dict = {}
|
||||
tensor_table_dict[feed_var_name] = {}
|
||||
tensor_table_dict[feed_var_name]["feed_var_name"] = feed_var_name
|
||||
tensor_table_dict[feed_var_name]["fetch_var_name"] = fetch_var_name
|
||||
tensor_table_dict[feed_var_name]["startup_program"] = startup_program
|
||||
tensor_table_dict[feed_var_name]["main_program"] = main_program
|
||||
tensor_table_dict[feed_var_name]["tensor_table_class"] = (
|
||||
tensor_table_class
|
||||
)
|
||||
attrs['tensor_table'] = tensor_table_dict
|
||||
|
||||
def _get_lr_scheduler_program(self, lr_scheduler, lr_decay_steps):
|
||||
scheduler_decay = [
|
||||
'NoamDecay',
|
||||
'NaturalExpDecay',
|
||||
'InverseTimeDecay',
|
||||
'ExponentialDecay',
|
||||
]
|
||||
|
||||
decay_main_program = paddle.static.Program()
|
||||
decay_startup_program = paddle.static.Program()
|
||||
lr_name = ""
|
||||
|
||||
if isinstance(lr_scheduler, ExponentialDecay):
|
||||
with paddle.static.program_guard(
|
||||
decay_main_program, decay_startup_program
|
||||
):
|
||||
lr = exponential_decay(
|
||||
1.0, lr_decay_steps, lr_scheduler.gamma, True
|
||||
)
|
||||
lr_name = lr.name
|
||||
logging.warning(
|
||||
f"ExponentialDecay is set, staircase = True, global learning rate decay step is [ {lr_decay_steps} ], Change decay steps as follow: \n"
|
||||
"\t strategy = paddle.distributed.fleet.DistributedStrategy() \n "
|
||||
"\t strategy.a_sync = True \n"
|
||||
"\t strategy.a_sync_configs= { 'lr_decay_steps' : YOUR_DECAY_STEP } \n"
|
||||
)
|
||||
elif isinstance(lr_scheduler, NoamDecay):
|
||||
with paddle.static.program_guard(
|
||||
decay_main_program, decay_startup_program
|
||||
):
|
||||
lr = noam_decay(
|
||||
lr_scheduler.d_model, lr_scheduler.warmup_steps, 1.0
|
||||
)
|
||||
lr_name = lr.name
|
||||
logging.warning(
|
||||
f"NoamDecay is set, warmup steps is [ {lr_scheduler.warmup_steps} ]"
|
||||
)
|
||||
elif isinstance(lr_scheduler, NaturalExpDecay):
|
||||
with paddle.static.program_guard(
|
||||
decay_main_program, decay_startup_program
|
||||
):
|
||||
lr = paddle.optimizer.lr.NaturalExpDecay(
|
||||
1.0, lr_scheduler.gamma
|
||||
).get_lr()
|
||||
lr_name = lr.name
|
||||
logging.warning(
|
||||
f"NaturalExpDecay is set, staircase = True, global learning rate decay step is [ {lr_decay_steps} ], Change decay steps as follow: \n"
|
||||
"\t strategy = paddle.distributed.fleet.DistributedStrategy() \n "
|
||||
"\t strategy.a_sync = True \n"
|
||||
"\t strategy.a_sync_configs= { 'lr_decay_steps' : YOUR_DECAY_STEP } \n"
|
||||
)
|
||||
elif isinstance(lr_scheduler, InverseTimeDecay):
|
||||
with paddle.static.program_guard(
|
||||
decay_main_program, decay_startup_program
|
||||
):
|
||||
lr = inverse_time_decay(
|
||||
1.0, lr_decay_steps, lr_scheduler.gamma, True
|
||||
)
|
||||
lr_name = lr.name
|
||||
logging.warning(
|
||||
f"InverseTimeDecay is set, staircase = True, global learning rate decay step is [ {lr_decay_steps} ], Change decay steps as follow: \n"
|
||||
"\t strategy = paddle.distributed.fleet.DistributedStrategy() \n "
|
||||
"\t strategy.a_sync = True \n"
|
||||
"\t strategy.a_sync_configs= { 'lr_decay_steps' : YOUR_DECAY_STEP } \n"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Not supported current LearningRate strategy, please use follow decay strategy: {scheduler_decay}"
|
||||
)
|
||||
|
||||
return decay_main_program, decay_startup_program, lr_name
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
attrs = pass_ctx._attrs
|
||||
if not hasattr(attrs['origin_main_program'], 'lr_scheduler'):
|
||||
return
|
||||
|
||||
assert isinstance(
|
||||
attrs['origin_main_program'].lr_scheduler, LRScheduler
|
||||
), "must be LRScheduler"
|
||||
|
||||
ops = get_optimize_ops(attrs['origin_main_program'])
|
||||
(
|
||||
lr_decay_main_program,
|
||||
lr_decay_startup_program,
|
||||
lr_name,
|
||||
) = self._get_lr_scheduler_program(
|
||||
attrs['origin_main_program'].lr_scheduler, attrs['lr_decay_steps']
|
||||
)
|
||||
self._add_tensor_table(
|
||||
attrs,
|
||||
"@LR_DECAY_COUNTER@",
|
||||
lr_name,
|
||||
lr_decay_startup_program,
|
||||
lr_decay_main_program,
|
||||
"GlobalStepTable",
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
@register_pass("add_listen_and_serv_pass")
|
||||
class AddListenAndServPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
attrs = pass_ctx._attrs
|
||||
opt = {
|
||||
"grad_to_block_id": None,
|
||||
"sparse_grad_to_param": None,
|
||||
"lr_decay_block_id": None,
|
||||
"dense_optimize_blocks": None,
|
||||
"sparse_optimize_blocks": None,
|
||||
# runtime attribute
|
||||
"endpoint": get_ps_endpoint(attrs['role_maker']),
|
||||
"pserver_id": get_role_id(attrs['role_maker']),
|
||||
"Fanin": get_trainers(attrs['role_maker']),
|
||||
"distributed_mode": attrs['ps_mode'],
|
||||
"rpc_get_thread_num": -1,
|
||||
"rpc_send_thread_num": -1,
|
||||
"rpc_prefetch_thread_num": -1,
|
||||
}
|
||||
main_program.global_block().append_op(
|
||||
type="listen_and_serv", inputs={'X': []}, outputs={}, attrs=opt
|
||||
)
|
||||
|
||||
|
||||
@register_pass("add_rpc_global_flags_pass")
|
||||
class AddRpcGlobalFlagsPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
pass
|
||||
|
||||
|
||||
@register_pass("add_optimizer_pass")
|
||||
class AddOptimizerPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
pass
|
||||
|
||||
|
||||
@register_pass("add_geo_optimizer_pass")
|
||||
class AddGeoOptimizerPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
pass
|
||||
|
||||
|
||||
@register_pass("build_pserver_startup_program_pass")
|
||||
class BuildPserverStartupProgramPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
pass
|
||||
|
||||
|
||||
@register_pass("delete_unused_in_startup_pass")
|
||||
class DeleteUnusedInStartupPass(PassBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def _check_self(self):
|
||||
return True
|
||||
|
||||
def _check_conflict(self, other_pass):
|
||||
return True
|
||||
|
||||
def _apply_single_impl(self, main_program, startup_program, pass_ctx):
|
||||
pass
|
||||
+1652
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user