chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,62 @@
|
||||
# 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 .base_cost import ( # noqa: F401
|
||||
CommContext,
|
||||
Cost,
|
||||
_g_op_cost_factory,
|
||||
build_comm_costs_from_descs,
|
||||
build_comm_desc,
|
||||
build_comm_desc_from_dist_op,
|
||||
build_comp_costs_from_descs,
|
||||
build_comp_desc_from_dist_op,
|
||||
build_comp_desc_str_for_predict,
|
||||
build_dp_costs,
|
||||
calc_time_by_cost_model,
|
||||
)
|
||||
from .comm_op_cost import ( # noqa: F401
|
||||
AllgatherOpCost,
|
||||
AllReduceOpCost,
|
||||
AllreduceSumOpCost,
|
||||
BroadcastOpCost,
|
||||
IdentityOpCost,
|
||||
RecvOpCost,
|
||||
SendOpCost,
|
||||
)
|
||||
from .comp_op_cost import ( # noqa: F401
|
||||
ConcatOpCost,
|
||||
EmbeddingGradOpCost,
|
||||
EmbeddingOpCost,
|
||||
FillConstantBatchSizeLikeOpCost,
|
||||
MatmulGradOpCost,
|
||||
MatmulOpCost,
|
||||
MatmulV2GradOpCost,
|
||||
MatmulV2OpCost,
|
||||
MulGradOpCost,
|
||||
MulOpCost,
|
||||
Reshape2GradOpCost,
|
||||
Reshape2OpCost,
|
||||
SliceOpCost,
|
||||
SoftmaxGradOpCost,
|
||||
SoftmaxOpCost,
|
||||
SplitOpCost,
|
||||
Transpose2GradOpCost,
|
||||
Transpose2OpCost,
|
||||
)
|
||||
from .estimate_cost import CostEstimator # noqa: F401
|
||||
from .op_runtime_cost import ( # noqa: F401
|
||||
check_if_op_supports_runtime_profiling,
|
||||
measure_program_real_op_cost,
|
||||
)
|
||||
from .tensor_cost import TensorCost # noqa: F401
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,316 @@
|
||||
# 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 math
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
|
||||
from .base_cost import CommOpCost, register_op_cost
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AllreduceSumOpCost(CommOpCost):
|
||||
OP_TYPE = "c_allreduce_sum"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
# use tree if cross machine and use ring if in a single machine
|
||||
time = None
|
||||
cluster = self.comm_context.cluster
|
||||
if not cluster.cross_machine(self.group_ranks):
|
||||
time = self.calc_time_ring()
|
||||
else:
|
||||
time = self.calc_time_tree()
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count - self.machine_count)
|
||||
* self.comm_context.intra_ring
|
||||
)
|
||||
alpha += (
|
||||
2
|
||||
* (self.machine_count - 1)
|
||||
* (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = (
|
||||
alpha
|
||||
+ 2
|
||||
* (self.rank_count - 1)
|
||||
/ self.rank_count
|
||||
* self.comm_count
|
||||
* beta
|
||||
)
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_tree(self):
|
||||
alpha = self.comm_context.base_tree
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count / self.machine_count - 1)
|
||||
* self.comm_context.intra_tree
|
||||
)
|
||||
alpha += math.log2(self.machine_count) * (
|
||||
self.comm_context.inter_tree + self.hops * self.comm_context.switch
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
|
||||
time = alpha + 2 * self.comm_count * beta
|
||||
|
||||
return time
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AllReduceOpCost(CommOpCost):
|
||||
OP_TYPE = "all_reduce"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
# use tree if cross machine and use ring if in a single machine
|
||||
time = None
|
||||
cluster = self.comm_context.cluster
|
||||
if not cluster.cross_machine(self.group_ranks):
|
||||
time = self.calc_time_ring()
|
||||
else:
|
||||
time = self.calc_time_tree()
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count - self.machine_count)
|
||||
* self.comm_context.intra_ring
|
||||
)
|
||||
alpha += (
|
||||
2
|
||||
* (self.machine_count - 1)
|
||||
* (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = (
|
||||
alpha
|
||||
+ 2
|
||||
* (self.rank_count - 1)
|
||||
/ self.rank_count
|
||||
* self.comm_count
|
||||
* beta
|
||||
)
|
||||
|
||||
return time
|
||||
|
||||
def calc_time_tree(self):
|
||||
alpha = self.comm_context.base_tree
|
||||
alpha += (
|
||||
2
|
||||
* (self.rank_count / self.machine_count - 1)
|
||||
* self.comm_context.intra_tree
|
||||
)
|
||||
alpha += math.log2(self.machine_count) * (
|
||||
self.comm_context.inter_tree + self.hops * self.comm_context.switch
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
|
||||
time = alpha + 2 * self.comm_count * beta
|
||||
|
||||
return time
|
||||
|
||||
@property
|
||||
def comm_count(self):
|
||||
from ..reshard import get_var_with_recursion
|
||||
|
||||
if self._comm_count is None:
|
||||
dtype = None
|
||||
shape = None
|
||||
if self.op is not None:
|
||||
vars = self.op.block.vars
|
||||
try:
|
||||
var_name = self.op.input("x")[0]
|
||||
except:
|
||||
var_name = self.op.output("out")[0]
|
||||
var = get_var_with_recursion(
|
||||
var_name, self.op.block, self.op.block.program
|
||||
)
|
||||
dtype = var.dtype
|
||||
shape = var.shape
|
||||
elif self.op_desc is not None:
|
||||
dtype = self.op_desc["inputs"]["x"][0][0]
|
||||
shape = self.op_desc["inputs"]["x"][0][1]
|
||||
|
||||
factor = None
|
||||
if dtype == paddle.float32 or dtype == paddle.int32:
|
||||
factor = 4
|
||||
else:
|
||||
raise ValueError(f"Unsupported comm dtype {dtype}")
|
||||
comm_count = int(np.prod(shape)) * factor
|
||||
self._comm_count = comm_count
|
||||
|
||||
return self._comm_count
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AllgatherOpCost(CommOpCost):
|
||||
OP_TYPE = "all_gather"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
time = self.calc_time_ring()
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
alpha += (
|
||||
self.rank_count - self.machine_count
|
||||
) * self.comm_context.intra_ring
|
||||
alpha += (self.machine_count - 1) * (
|
||||
self.comm_context.inter_ring + self.hops * self.comm_context.switch
|
||||
)
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = (
|
||||
alpha
|
||||
+ (self.rank_count - 1) / self.rank_count * self.comm_count * beta
|
||||
)
|
||||
return time
|
||||
|
||||
@property
|
||||
def comm_count(self):
|
||||
from ..reshard import get_var_with_recursion
|
||||
|
||||
if self._comm_count is None:
|
||||
dtype = None
|
||||
shape = None
|
||||
if self.op is not None:
|
||||
vars = self.op.block.vars
|
||||
try:
|
||||
var_name = self.op.input("x")[0]
|
||||
except:
|
||||
var_name = self.op.output("out")[0]
|
||||
var = get_var_with_recursion(
|
||||
var_name, self.op.block, self.op.block.program
|
||||
)
|
||||
dtype = var.dtype
|
||||
shape = var.shape
|
||||
elif self.op_desc is not None:
|
||||
dtype = self.op_desc["inputs"]["X"][0][0]
|
||||
shape = self.op_desc["inputs"]["X"][0][1]
|
||||
|
||||
factor = None
|
||||
if dtype == paddle.float32 or dtype == paddle.int32:
|
||||
factor = 4
|
||||
else:
|
||||
raise ValueError(f"Unsupported comm dtype {dtype}")
|
||||
comm_count = int(np.prod(shape)) * factor
|
||||
self._comm_count = comm_count
|
||||
|
||||
return self._comm_count
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class BroadcastOpCost(CommOpCost):
|
||||
OP_TYPE = "broadcast"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
time = self.calc_time_ring()
|
||||
return time
|
||||
|
||||
def calc_time_ring(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
if self.machine_count > 1:
|
||||
alpha += (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
else:
|
||||
alpha += self.comm_context.intra_ring
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = alpha + self.comm_count * beta
|
||||
|
||||
return time
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class IdentityOpCost(CommOpCost):
|
||||
OP_TYPE = "c_identity"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
return self.comm_count * 1 / (144 * 1e3)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class RecvOpCost(CommOpCost):
|
||||
OP_TYPE = "recv_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
if self.machine_count > 1:
|
||||
alpha += (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
else:
|
||||
alpha += self.comm_context.intra_ring
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = alpha + self.comm_count * beta
|
||||
return time
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SendOpCost(CommOpCost):
|
||||
OP_TYPE = "send_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, comm_context=None):
|
||||
super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)
|
||||
|
||||
def calc_time(self):
|
||||
alpha = self.comm_context.base_ring
|
||||
if self.machine_count > 1:
|
||||
alpha += (
|
||||
self.comm_context.inter_ring
|
||||
+ self.hops * self.comm_context.switch
|
||||
)
|
||||
else:
|
||||
alpha += self.comm_context.intra_ring
|
||||
beta = self.comm_context.get_max_beta(self.group_ranks)
|
||||
time = alpha + self.comm_count * beta
|
||||
|
||||
return time
|
||||
@@ -0,0 +1,591 @@
|
||||
# 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 .base_cost import CompOpCost, register_op_cost
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AdamOpCost(CompOpCost):
|
||||
OP_TYPE = "adam"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ArgsortOpCost(CompOpCost):
|
||||
OP_TYPE = "argsort"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AssignOpCost(CompOpCost):
|
||||
OP_TYPE = "assign"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class AssignValueOpCost(CompOpCost):
|
||||
OP_TYPE = "assign_value"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class BeamSearchOpCost(CompOpCost):
|
||||
OP_TYPE = "beam_search"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class BeamSearchDecodeOpCost(CompOpCost):
|
||||
OP_TYPE = "beam_search_decode"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class CastOpCost(CompOpCost):
|
||||
OP_TYPE = "cast"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ConcatOpCost(CompOpCost):
|
||||
OP_TYPE = "concat"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class DropoutOpCost(CompOpCost):
|
||||
OP_TYPE = "dropout"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class DropoutGradOpCost(CompOpCost):
|
||||
OP_TYPE = "dropout_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseAddOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_add"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseAddGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_add_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseDivOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_div"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseDivGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_div_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseMulOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_mul"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseMulGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_mul_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseSubOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_sub"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ElementwiseSubGradOpCost(CompOpCost):
|
||||
OP_TYPE = "elementwise_sub_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class EqualOpCost(CompOpCost):
|
||||
OP_TYPE = "equal"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class EmbeddingOpCost(CompOpCost):
|
||||
OP_TYPE = "c_embedding"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class EmbeddingGradOpCost(CompOpCost):
|
||||
OP_TYPE = "c_embedding_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FillConstantOpCost(CompOpCost):
|
||||
OP_TYPE = "fill_constant"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FillConstantBatchSizeLikeOpCost(CompOpCost):
|
||||
OP_TYPE = "fill_constant_batch_size_like"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FusedSoftmaxMaskUpperTriangleOpCost(CompOpCost):
|
||||
OP_TYPE = "fused_softmax_mask_upper_triangle"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class FusedSoftmaxMaskUpperTriangleGradOpCost(CompOpCost):
|
||||
OP_TYPE = "fused_softmax_mask_upper_triangle_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GatherOpCost(CompOpCost):
|
||||
OP_TYPE = "gather"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GeluOpCost(CompOpCost):
|
||||
OP_TYPE = "gelu"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GeluGradOpCost(CompOpCost):
|
||||
OP_TYPE = "gelu_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class GreaterEqualOpCost(CompOpCost):
|
||||
OP_TYPE = "greater_equal"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class IncrementOpCost(CompOpCost):
|
||||
OP_TYPE = "increment"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class IsEmptyOpCost(CompOpCost):
|
||||
OP_TYPE = "is_empty"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LayerNormOpCost(CompOpCost):
|
||||
OP_TYPE = "layer_norm"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LayerNormGradOpCost(CompOpCost):
|
||||
OP_TYPE = "layer_norm_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LessThanOpCost(CompOpCost):
|
||||
OP_TYPE = "less_than"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LogicalNotOpCost(CompOpCost):
|
||||
OP_TYPE = "logical_not"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LogicalAndOpCost(CompOpCost):
|
||||
OP_TYPE = "logical_and"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LodResetOpCost(CompOpCost):
|
||||
OP_TYPE = "lod_reset"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LogOpCost(CompOpCost):
|
||||
OP_TYPE = "log"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LookupTableV2OpCost(CompOpCost):
|
||||
OP_TYPE = "lookup_table_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class LookupTableV2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "lookup_table_v2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulOpCost(CompOpCost):
|
||||
OP_TYPE = "matmul"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulGradOpCost(CompOpCost):
|
||||
OP_TYPE = "matmul_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulV2OpCost(CompOpCost):
|
||||
OP_TYPE = "matmul_v2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MatmulV2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "matmul_v2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MemcpyOpCost(CompOpCost):
|
||||
OP_TYPE = "memcpy"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MulOpCost(CompOpCost):
|
||||
OP_TYPE = "mul"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class MulGradOpCost(CompOpCost):
|
||||
OP_TYPE = "mul_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class OneHotOpCost(CompOpCost):
|
||||
OP_TYPE = "one_hot"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReadFromArrayOpCost(CompOpCost):
|
||||
OP_TYPE = "read_from_array"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceSumOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_sum"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceSumGradOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_sum_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Reshape2OpCost(CompOpCost):
|
||||
OP_TYPE = "reshape2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Reshape2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "reshape2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceMeanOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_mean"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ReduceMeanGradOpCost(CompOpCost):
|
||||
OP_TYPE = "reduce_mean_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ScaleOpCost(CompOpCost):
|
||||
OP_TYPE = "scale"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class ShapeOpCost(CompOpCost):
|
||||
OP_TYPE = "shape"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SliceOpCost(CompOpCost):
|
||||
OP_TYPE = "slice"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxGradOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxWithCrossEntropyOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax_with_cross_entropy"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SoftmaxWithCrossEntropyGradOpCost(CompOpCost):
|
||||
OP_TYPE = "softmax_with_cross_entropy_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SplitOpCost(CompOpCost):
|
||||
OP_TYPE = "split"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Squeeze2OpCost(CompOpCost):
|
||||
OP_TYPE = "squeeze2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SquareOpCost(CompOpCost):
|
||||
OP_TYPE = "square"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SquareGradOpCost(CompOpCost):
|
||||
OP_TYPE = "square_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class SumOpCost(CompOpCost):
|
||||
OP_TYPE = "sum"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class TopKOpCost(CompOpCost):
|
||||
OP_TYPE = "top_k"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Transpose2OpCost(CompOpCost):
|
||||
OP_TYPE = "transpose2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Transpose2GradOpCost(CompOpCost):
|
||||
OP_TYPE = "transpose2_grad"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class Unsqueeze2OpCost(CompOpCost):
|
||||
OP_TYPE = "unsqueeze2"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
|
||||
|
||||
@register_op_cost
|
||||
class WriteToArrayOpCost(CompOpCost):
|
||||
OP_TYPE = "write_to_array"
|
||||
|
||||
def __init__(self, op=None, op_desc=None, cluster=None, rank=None):
|
||||
super().__init__(op=op, op_desc=op_desc, cluster=cluster, rank=rank)
|
||||
@@ -0,0 +1,671 @@
|
||||
# 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
|
||||
from functools import reduce
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
||||
|
||||
from ..dist_tensor import DistributedTensor
|
||||
from ..operators.common import get_distributed_operator_impl_container
|
||||
from .base_cost import Cost
|
||||
|
||||
|
||||
class CostEstimator:
|
||||
_special_op_type = ["fused_attention", "fused_feedforward"]
|
||||
|
||||
def __init__(
|
||||
self, program, cluster, mode="modeling", rank=None, loop_count=10
|
||||
):
|
||||
self._program = program
|
||||
self._cluster = cluster
|
||||
self._check_mode(mode)
|
||||
self._mode = mode
|
||||
self._rank = rank if rank is not None else paddle.distributed.get_rank()
|
||||
self._loop_count = loop_count
|
||||
self._global_cost = Cost()
|
||||
self._local_cost_mapping = {}
|
||||
self._detailed_cost = OrderedDict() # {`op_id`: {"reshard": [], "dist_op": [], "local_cost": local_cost}}}
|
||||
self._bubble_time_mapping = {}
|
||||
self._ordered_ops = []
|
||||
self.max_memories = {}
|
||||
self.max_memory = None
|
||||
|
||||
@property
|
||||
def loop_count(self):
|
||||
return self._loop_count
|
||||
|
||||
@property
|
||||
def detailed_cost(self):
|
||||
return self._detailed_cost
|
||||
|
||||
@property
|
||||
def program(self):
|
||||
return self._program
|
||||
|
||||
@property
|
||||
def rank(self):
|
||||
return self._rank
|
||||
|
||||
@property
|
||||
def dist_context(self):
|
||||
return self._dist_context
|
||||
|
||||
@property
|
||||
def cluster(self):
|
||||
return self._cluster
|
||||
|
||||
@property
|
||||
def mode(self):
|
||||
return self._mode
|
||||
|
||||
@property
|
||||
def global_cost(self):
|
||||
max_time = 0
|
||||
memory = 0
|
||||
flops = 0
|
||||
for rank in self._local_cost_mapping:
|
||||
cost = self._local_cost_mapping[rank]
|
||||
if cost.time > max_time:
|
||||
max_time = cost.time
|
||||
memory += cost.memory
|
||||
flops += cost.flops
|
||||
self._global_cost.time = max_time
|
||||
self._global_cost.memory = memory
|
||||
self._global_cost.flops = flops
|
||||
return self._global_cost
|
||||
|
||||
def local_cost(self, rank=None):
|
||||
rank = self.rank if rank is None else rank
|
||||
if rank not in self._local_cost_mapping:
|
||||
self._local_cost_mapping[rank] = Cost()
|
||||
|
||||
return self._local_cost_mapping[rank]
|
||||
|
||||
def local_bubble_time(self, rank=None):
|
||||
rank = self.rank if rank is None else rank
|
||||
return self._bubble_time_mapping[rank]
|
||||
|
||||
def _check_mode(self, mode):
|
||||
if mode not in ["modeling", "profiling"]:
|
||||
raise ValueError(
|
||||
f"Just support modeling and profiling, but got {mode}"
|
||||
)
|
||||
|
||||
def _is_special_var_name(self, var_name):
|
||||
special_var_name = ["lod_tensor_blocking_queue_0"]
|
||||
if var_name in special_var_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _estimate_core(self, dist_context, resharder, block):
|
||||
from ..reshard import get_var_with_recursion
|
||||
|
||||
ops = block.ops
|
||||
loop_count = None
|
||||
if block.desc.id != self.program.global_block().desc.id:
|
||||
loop_count = self.loop_count
|
||||
else:
|
||||
loop_count = 1
|
||||
for i in range(loop_count):
|
||||
for op in ops:
|
||||
self._detailed_cost[op.desc.id()] = OrderedDict()
|
||||
# If in the while sub block, the detail of cost is the last cost
|
||||
detail = self._detailed_cost[op.desc.id()]
|
||||
detail["reshard_cost"] = OrderedDict() #
|
||||
detail["dist_op_cost"] = []
|
||||
if int(op.attr('op_role')) == int(OpRole.Optimize):
|
||||
continue
|
||||
if op.type in [
|
||||
"create_py_reader",
|
||||
"create_double_buffer_reader",
|
||||
"read",
|
||||
]:
|
||||
continue
|
||||
|
||||
# NOTE: It does not support nested loop and just supports while op when op has sub block now.
|
||||
if op.type == "while":
|
||||
while_block = self.program.blocks[op.attr("sub_block").id]
|
||||
self._estimate_core(dist_context, resharder, while_block)
|
||||
continue
|
||||
|
||||
for var_name in op.input_arg_names:
|
||||
if self._is_special_var_name(var_name):
|
||||
continue
|
||||
var = get_var_with_recursion(var_name, block, self.program)
|
||||
reshard_cost = resharder.get_cost(op, var, self.cluster)
|
||||
|
||||
# Calc reshard cost
|
||||
if reshard_cost is not None:
|
||||
detail["reshard_cost"][var_name] = reshard_cost
|
||||
|
||||
comm_costs = reshard_cost[0]
|
||||
local_comp_cost = reshard_cost[1]
|
||||
for comm_cost in comm_costs:
|
||||
# Time is cumulative in global cost and local cost, but memory and flops just are cumulative in global cost.
|
||||
# Comm sync
|
||||
for item in comm_cost:
|
||||
group_ranks, cost = item
|
||||
max_time = None
|
||||
cost_time = {}
|
||||
for rank in group_ranks:
|
||||
rank_cost = self.local_cost(rank)
|
||||
cost_time[rank] = rank_cost.time
|
||||
if max_time is None:
|
||||
max_time = rank_cost.time
|
||||
else:
|
||||
if max_time < rank_cost.time:
|
||||
max_time = rank_cost.time
|
||||
|
||||
for rank in group_ranks:
|
||||
self.local_cost(rank).time = (
|
||||
max_time + cost.time
|
||||
)
|
||||
|
||||
if rank not in self._bubble_time_mapping:
|
||||
self._bubble_time_mapping[rank] = 0
|
||||
|
||||
self._bubble_time_mapping[rank] += (
|
||||
max_time - cost_time[rank]
|
||||
)
|
||||
|
||||
for rank in local_comp_cost:
|
||||
for comp_cost in local_comp_cost[rank]:
|
||||
self.local_cost(rank).time += comp_cost.time
|
||||
|
||||
# Calc dist op cost
|
||||
dist_op = dist_context.get_dist_op_for_program(op)
|
||||
if not dist_op:
|
||||
continue
|
||||
|
||||
op_dist_attr = dist_op.dist_attr
|
||||
processes = op_dist_attr.process_mesh.process_ids
|
||||
|
||||
container = get_distributed_operator_impl_container(
|
||||
op_dist_attr.impl_type
|
||||
)
|
||||
dist_impl = container.impls[op_dist_attr.impl_idx]
|
||||
|
||||
dist_op_cost = dist_impl.calc_cost(
|
||||
op.attr('op_role'), dist_op, dist_context, self.cluster
|
||||
)
|
||||
detail["dist_op_cost"] = dist_op_cost
|
||||
|
||||
if dist_op_cost is None:
|
||||
assert (
|
||||
dist_op.serial_op.type in CostEstimator._special_op_type
|
||||
)
|
||||
continue
|
||||
for item in dist_op_cost:
|
||||
if isinstance(item, list):
|
||||
# Comm sync
|
||||
for comm_op_cost in item:
|
||||
max_time = None
|
||||
cost_time = {}
|
||||
group_ranks = comm_op_cost.group_ranks
|
||||
for rank in comm_op_cost.group_ranks:
|
||||
rank_cost = self.local_cost(rank)
|
||||
cost_time[rank] = rank_cost.time
|
||||
if max_time is None:
|
||||
max_time = rank_cost.time
|
||||
else:
|
||||
if max_time < rank_cost.time:
|
||||
max_time = rank_cost.time
|
||||
for rank in group_ranks:
|
||||
self.local_cost(rank).time = (
|
||||
max_time + comm_op_cost.time
|
||||
if op.attr('op_role') != OpRole.Backward
|
||||
else max_time + 0.9 * comm_op_cost.time
|
||||
)
|
||||
if rank not in self._bubble_time_mapping:
|
||||
self._bubble_time_mapping[rank] = 0
|
||||
self._bubble_time_mapping[rank] += (
|
||||
max_time - cost_time[rank]
|
||||
)
|
||||
elif isinstance(item, dict):
|
||||
# Op just one
|
||||
for rank in processes:
|
||||
# DP+PP+MP
|
||||
if rank not in item:
|
||||
continue
|
||||
self.local_cost(rank).time += item[rank].time
|
||||
|
||||
def prepare(self):
|
||||
self._global_cost = Cost()
|
||||
self._local_cost_mapping = {}
|
||||
self._detailed_cost = OrderedDict()
|
||||
self._bubble_time_mapping = {}
|
||||
|
||||
def _calculate_bytes(self, sizes, dtype):
|
||||
if sizes:
|
||||
total_count = reduce(lambda x, y: x * y, sizes, 1)
|
||||
else:
|
||||
total_count = 0
|
||||
|
||||
if dtype == paddle.float64 or dtype == paddle.int64:
|
||||
dtype_factor = 8
|
||||
elif dtype == paddle.float32 or dtype == paddle.int32:
|
||||
dtype_factor = 4
|
||||
elif (
|
||||
dtype == paddle.float16
|
||||
or dtype == paddle.bfloat16
|
||||
or dtype == paddle.int16
|
||||
):
|
||||
dtype_factor = 2
|
||||
elif dtype == paddle.int8 or dtype == paddle.uint8:
|
||||
dtype_factor = 1
|
||||
else:
|
||||
dtype_factor = 8
|
||||
|
||||
memory = total_count * dtype_factor
|
||||
return memory
|
||||
|
||||
def _estimate_max_memory_by_dist_op(self, dist_context):
|
||||
# This estimation will be improved, now reshard and inplace are not considered.
|
||||
# Persist var is not free.
|
||||
def _convert_pm_and_dm_to_str(process_mesh, dims_mapping):
|
||||
processes = ",".join([str(x) for x in process_mesh.process_ids])
|
||||
topology = ",".join([str(x) for x in process_mesh.shape])
|
||||
dims_mapping = ",".join([str(x) for x in dims_mapping])
|
||||
result = processes + topology + dims_mapping
|
||||
return result
|
||||
|
||||
memories = {}
|
||||
self.max_memories = {}
|
||||
var_info = {} # var_name: [[process_mesh, dims_mapping], [id]], [[process_mesh, dims_mapping], [id]]}
|
||||
|
||||
for block in self.program.blocks:
|
||||
for op in block.ops:
|
||||
self._ordered_ops.append([op.desc.id(), op])
|
||||
self._ordered_ops.sort(key=lambda x: x[0])
|
||||
|
||||
parameters = set()
|
||||
for op_id, op in self._ordered_ops:
|
||||
if op.type in [
|
||||
"create_py_reader",
|
||||
"create_double_buffer_reader",
|
||||
"read",
|
||||
]:
|
||||
continue
|
||||
dist_op = dist_context.get_dist_op_for_program(op)
|
||||
if not dist_op:
|
||||
continue
|
||||
process_mesh = dist_op.dist_attr.process_mesh
|
||||
for var_name in op.input_arg_names:
|
||||
input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
|
||||
if var_name not in var_info:
|
||||
var_info[var_name] = {}
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, input_dims_mapping
|
||||
)
|
||||
if key not in var_info[var_name]:
|
||||
var_info[var_name][key] = {}
|
||||
# It is even partition now
|
||||
if "position" not in var_info[var_name][key]:
|
||||
var_info[var_name][key]["position"] = []
|
||||
var_info[var_name][key]["position"].append(op_id)
|
||||
|
||||
if "memory" not in var_info[var_name][key]:
|
||||
var = dist_op.get_serial_input(var_name)
|
||||
global_sizes = var.shape
|
||||
dtype = var.dtype
|
||||
sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes,
|
||||
input_dims_mapping,
|
||||
process_mesh.shape,
|
||||
process_mesh.process_ids,
|
||||
)
|
||||
var_info[var_name][key]["memory"] = self._calculate_bytes(
|
||||
sizes, dtype
|
||||
)
|
||||
if var.persistable:
|
||||
name = var_name + key
|
||||
if name not in parameters:
|
||||
parameters.add(name)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key][
|
||||
"memory"
|
||||
]
|
||||
|
||||
for var_name in op.output_arg_names:
|
||||
output_dims_mapping = dist_op.dist_attr.get_output_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
if var_name not in var_info:
|
||||
var_info[var_name] = {}
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, output_dims_mapping
|
||||
)
|
||||
if key not in var_info[var_name]:
|
||||
var_info[var_name][key] = {}
|
||||
if "position" not in var_info[var_name][key]:
|
||||
var_info[var_name][key]["position"] = []
|
||||
var_info[var_name][key]["position"].append(op_id)
|
||||
|
||||
if "memory" not in var_info[var_name][key]:
|
||||
var = dist_op.get_serial_output(var_name)
|
||||
global_sizes = var.shape
|
||||
dtype = var.dtype
|
||||
sizes = DistributedTensor.get_local_sizes(
|
||||
global_sizes,
|
||||
output_dims_mapping,
|
||||
process_mesh.shape,
|
||||
process_mesh.process_ids,
|
||||
)
|
||||
var_info[var_name][key]["memory"] = self._calculate_bytes(
|
||||
sizes, dtype
|
||||
)
|
||||
if var.persistable:
|
||||
name = var_name + key
|
||||
if name not in parameters:
|
||||
parameters.add(name)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key][
|
||||
"memory"
|
||||
]
|
||||
|
||||
has_used_vars = set()
|
||||
not_calc_vars = set()
|
||||
for op_id, op in self._ordered_ops:
|
||||
if op.type in [
|
||||
"create_py_reader",
|
||||
"create_double_buffer_reader",
|
||||
"read",
|
||||
]:
|
||||
continue
|
||||
can_free_memories = {}
|
||||
can_free_vars = set()
|
||||
dist_op = dist_context.get_dist_op_for_program(op)
|
||||
if not dist_op:
|
||||
continue
|
||||
process_mesh = dist_op.dist_attr.process_mesh
|
||||
for var_name in op.input_arg_names:
|
||||
input_dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, input_dims_mapping
|
||||
)
|
||||
has_used_var = var_name + key
|
||||
var = dist_op.get_serial_input(var_name)
|
||||
# Not used
|
||||
if (
|
||||
has_used_var not in has_used_vars
|
||||
and has_used_var not in parameters
|
||||
):
|
||||
if has_used_var in not_calc_vars:
|
||||
continue
|
||||
has_used_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key]["memory"]
|
||||
# Used
|
||||
if op_id == var_info[var_name][key]["position"][-1]:
|
||||
if (
|
||||
has_used_var not in can_free_vars
|
||||
and not var.persistable
|
||||
):
|
||||
can_free_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in can_free_memories:
|
||||
can_free_memories[process] = 0
|
||||
can_free_memories[process] += var_info[var_name][
|
||||
key
|
||||
]["memory"]
|
||||
|
||||
for var_name in op.output_arg_names:
|
||||
output_dims_mapping = dist_op.dist_attr.get_output_dims_mapping(
|
||||
var_name
|
||||
)
|
||||
key = _convert_pm_and_dm_to_str(
|
||||
process_mesh, output_dims_mapping
|
||||
)
|
||||
has_used_var = var_name + key
|
||||
var = dist_op.get_serial_output(var_name)
|
||||
if (
|
||||
op.type == "reshape2"
|
||||
or op.type == "transpose2"
|
||||
or op.type == "elementwise_add"
|
||||
):
|
||||
not_calc_vars.add(has_used_var)
|
||||
continue
|
||||
# Not used
|
||||
if (
|
||||
has_used_var not in has_used_vars
|
||||
and has_used_var not in parameters
|
||||
):
|
||||
has_used_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in memories:
|
||||
memories[process] = 0
|
||||
memories[process] += var_info[var_name][key]["memory"]
|
||||
# Used
|
||||
if op_id == var_info[var_name][key]["position"][-1]:
|
||||
if (
|
||||
has_used_var not in can_free_vars
|
||||
and not var.persistable
|
||||
):
|
||||
can_free_vars.add(has_used_var)
|
||||
for process in process_mesh.process_ids:
|
||||
if process not in can_free_memories:
|
||||
can_free_memories[process] = 0
|
||||
can_free_memories[process] += var_info[var_name][
|
||||
key
|
||||
]["memory"]
|
||||
|
||||
# Calc peak memory
|
||||
for process in memories:
|
||||
if process not in self.max_memories:
|
||||
self.max_memories[process] = memories[process]
|
||||
else:
|
||||
if memories[process] > self.max_memories[process]:
|
||||
self.max_memories[process] = memories[process]
|
||||
# Free memory
|
||||
for process in can_free_memories:
|
||||
if process in memories:
|
||||
memories[process] -= can_free_memories[process]
|
||||
|
||||
# Calculate the max memory in all ranks
|
||||
max_memory = max(self.max_memories.values())
|
||||
self.max_memory = max_memory
|
||||
|
||||
return max_memory
|
||||
|
||||
def estimate(self, dist_context, resharder=None):
|
||||
self.prepare()
|
||||
from ..reshard import Resharder
|
||||
|
||||
resharder = (
|
||||
Resharder(self.program, None, self.rank, dist_context, [])
|
||||
if resharder is None
|
||||
else resharder
|
||||
)
|
||||
|
||||
block = self.program.global_block()
|
||||
self._estimate_core(dist_context, resharder, block)
|
||||
|
||||
return self.global_cost
|
||||
|
||||
def _print_tag(self, max_len, length):
|
||||
tag = "+" + "-" * max_len
|
||||
for i in range(length):
|
||||
print(tag, end="")
|
||||
if i == length - 1:
|
||||
print("+")
|
||||
|
||||
def _print_vals(self, vals, max_len):
|
||||
for idx, val in enumerate(vals):
|
||||
s = "|" + str(val).center(max_len)
|
||||
print(s, end="")
|
||||
if idx == len(vals) - 1:
|
||||
print("|")
|
||||
|
||||
def _pretty_print_memory_cost(self):
|
||||
"""Print memory of every rank prettily."""
|
||||
if not self.max_memories or not self.max_memory:
|
||||
raise ValueError("Please calculate memory cost before print.")
|
||||
|
||||
# Padding automatically
|
||||
max_len = 0
|
||||
header = ["Rank", "Memory(MiB)"]
|
||||
memories = [
|
||||
int(item // 1e6) for item in list(self.max_memories.values())
|
||||
]
|
||||
for memory in memories + header:
|
||||
if len(str(memory)) > max_len:
|
||||
max_len = len(str(memory))
|
||||
max_len += 4 # for pretty print of center
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print header
|
||||
self._print_vals(header, max_len)
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print rank and its memory
|
||||
for i in range(len(self.max_memories)):
|
||||
memory = memories[i]
|
||||
vals = [i, memory]
|
||||
self._print_vals(vals, max_len)
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
def _pretty_print_global(self):
|
||||
"""Print global execution time and max memory prettily."""
|
||||
if not self.max_memories or not self.max_memory:
|
||||
raise ValueError("Please calculate cost before print.")
|
||||
|
||||
# Padding automatically
|
||||
max_len = 0
|
||||
header = ["Execution Time(us)", "Max Memory(MiB)"]
|
||||
vals = [round(self.global_cost.time, 3), int(self.max_memory // 1e6)]
|
||||
for memory in vals + header:
|
||||
if len(str(memory)) > max_len:
|
||||
max_len = len(str(memory))
|
||||
max_len += 4 # for pretty print of center
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print header
|
||||
self._print_vals(header, max_len)
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
# Print exec time and max memory
|
||||
self._print_vals(vals, max_len)
|
||||
|
||||
# Print tag
|
||||
self._print_tag(max_len, len(header))
|
||||
|
||||
def pretty_print_cost(self):
|
||||
"""Print cost prettily."""
|
||||
print("The global execution time and max memory are as follows:")
|
||||
self._pretty_print_global()
|
||||
print("The memory of every rank is as follows:")
|
||||
self._pretty_print_memory_cost()
|
||||
|
||||
|
||||
def get_cost_from_engine(engine, mode):
|
||||
import copy
|
||||
|
||||
from ..utils import to_list
|
||||
|
||||
# Construct cost estimator by original main program
|
||||
serial_main_prog = (
|
||||
engine._fwd_main_progs[mode].clone()
|
||||
if mode in engine._fwd_main_progs
|
||||
else engine._orig_main_prog.clone()
|
||||
)
|
||||
|
||||
serial_startup_prog = (
|
||||
engine._fwd_dist_contexts[mode]._original_serial_main_program.clone()
|
||||
if mode in engine._fwd_dist_contexts
|
||||
else engine._orig_startup_prog.clone()
|
||||
)
|
||||
losses = (
|
||||
to_list(engine._loss)
|
||||
if (
|
||||
not isinstance(engine._loss, paddle.nn.Layer)
|
||||
and not callable(engine._loss)
|
||||
)
|
||||
else engine._losses
|
||||
)
|
||||
serial_optimizer = copy.deepcopy(engine._orig_optimizer)
|
||||
if mode in engine._fwd_dist_contexts:
|
||||
dist_context = copy.deepcopy(engine._fwd_dist_contexts[mode])
|
||||
else:
|
||||
from ..dist_context import DistributedContext
|
||||
|
||||
dist_context = DistributedContext(
|
||||
serial_main_prog,
|
||||
serial_startup_prog,
|
||||
serial_optimizer,
|
||||
losses,
|
||||
{},
|
||||
{"loss": losses},
|
||||
engine._cluster,
|
||||
engine._strategy,
|
||||
)
|
||||
from ..completion import Completer
|
||||
|
||||
completer = Completer(dist_context)
|
||||
completer.complete_forward_annotation()
|
||||
dist_context.block_state.parse_forward_blocks(
|
||||
dist_context.serial_main_program
|
||||
)
|
||||
|
||||
if mode == "eval" or mode == "predict":
|
||||
cost_estimator = CostEstimator(serial_main_prog, engine._cluster)
|
||||
elif mode == "train":
|
||||
from ..parallelizer_v2 import Parallelizer
|
||||
|
||||
# Get serial main program with backward
|
||||
parallelizer = Parallelizer(mode, completer, dist_context)
|
||||
# Generate backward
|
||||
loss_name = dist_context.serial_loss.name
|
||||
serial_loss = serial_main_prog.global_block()._var_recursive(loss_name)
|
||||
params_grads = parallelizer._generate_backward(
|
||||
serial_main_prog, serial_startup_prog, serial_loss
|
||||
)
|
||||
|
||||
# Generate optimizer
|
||||
optimizer_ops = parallelizer._generate_optimizer(
|
||||
serial_main_prog,
|
||||
serial_startup_prog,
|
||||
serial_optimizer,
|
||||
params_grads,
|
||||
)
|
||||
cost_estimator = CostEstimator(serial_main_prog, engine._cluster)
|
||||
|
||||
# Estimate global_cost and max memory
|
||||
global_cost = cost_estimator.estimate(dist_context)
|
||||
max_memory = cost_estimator._estimate_max_memory_by_dist_op(dist_context)
|
||||
|
||||
# Print the cost
|
||||
cost_estimator.pretty_print_cost()
|
||||
|
||||
return global_cost, max_memory
|
||||
@@ -0,0 +1,320 @@
|
||||
# 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 warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
from paddle.base.data_feeder import convert_dtype
|
||||
from paddle.base.executor import (
|
||||
_as_lodtensor,
|
||||
_StandaloneExecutor,
|
||||
check_feed_shape_type,
|
||||
)
|
||||
from paddle.base.framework import Operator, Program
|
||||
from paddle.distributed.auto_parallel.static.utils import get_logger, is_comm_op
|
||||
|
||||
|
||||
def check_if_op_supports_runtime_profiling(op):
|
||||
return not is_comm_op(op)
|
||||
|
||||
|
||||
def _measure_program_real_op_cost_multipass(program, place, run_iters, verbose):
|
||||
'''
|
||||
Run op profiling for a single pass. Internal function, do not call this directly.
|
||||
'''
|
||||
|
||||
# clone the program to avoid accidental change made to the vanilla program.
|
||||
cloned_program = program.clone()
|
||||
cloned_main_block = cloned_program.global_block()
|
||||
|
||||
# We will run the executor in a newly created scope, so that our
|
||||
# executor will not pollute the global scope when running. Since
|
||||
# we created a brand new scope, we need to manually create input
|
||||
# tensors and network parameters and feed fake data into them.
|
||||
scope = core.Scope()
|
||||
|
||||
logger = get_logger(log_level=logging.INFO)
|
||||
|
||||
def _analyze_graph_and_collect_all_vars_with_zero_in_degree():
|
||||
var_in_degree = {}
|
||||
|
||||
def _collect_op_input_var_names(op: Operator):
|
||||
input_var_names = []
|
||||
for input_name in op.input_names:
|
||||
input_var_names += op.input(input_name)
|
||||
return input_var_names
|
||||
|
||||
def _collect_op_output_var_names(op: Operator):
|
||||
output_var_names = []
|
||||
for output_name in op.output_names:
|
||||
output_var_names += op.output(output_name)
|
||||
return output_var_names
|
||||
|
||||
def _record_op_output_vars_in_degree(in_var_names, out_var_names):
|
||||
for out_var_name in out_var_names:
|
||||
if out_var_name in in_var_names:
|
||||
# NOTE (liuchenghao): if an op's input var is its output var,
|
||||
# this means this var forms an in-place connection to itself,
|
||||
# in this situation we need to ignore this variable, this way
|
||||
# we can ensure that vars with zero in-degree are dangling vars
|
||||
# and they should be created manually before program executes.
|
||||
continue
|
||||
var_in_degree[out_var_name] += 1
|
||||
|
||||
def _filter_vars_with_zero_in_degree_and_ignore_feed_fetch_vars():
|
||||
filtered_vars = []
|
||||
for var_name in var_in_degree:
|
||||
if var_name in ['feed', 'fetch']:
|
||||
continue
|
||||
if var_in_degree[var_name] == 0:
|
||||
filtered_vars.append(var_name)
|
||||
return filtered_vars
|
||||
|
||||
for op in cloned_main_block.ops:
|
||||
op: Operator
|
||||
if is_comm_op(op):
|
||||
# ignore communication op from graph, because sometimes we want to profile a sub-graph
|
||||
# and these dangling operators will not work (no graph to communicate to/from)
|
||||
continue
|
||||
input_var_names, output_var_names = (
|
||||
_collect_op_input_var_names(op),
|
||||
_collect_op_output_var_names(op),
|
||||
)
|
||||
for var_name in input_var_names + output_var_names:
|
||||
if var_name not in var_in_degree:
|
||||
var_in_degree[var_name] = 0
|
||||
_record_op_output_vars_in_degree(input_var_names, output_var_names)
|
||||
return _filter_vars_with_zero_in_degree_and_ignore_feed_fetch_vars()
|
||||
|
||||
def _alloc_and_fill_var(var_name):
|
||||
supported_var_dtypes = [
|
||||
"paddle.float16",
|
||||
"paddle.float32",
|
||||
"paddle.float64",
|
||||
"paddle.int8",
|
||||
"paddle.int16",
|
||||
"paddle.int32",
|
||||
"paddle.int64",
|
||||
"paddle.bool",
|
||||
]
|
||||
var = cloned_main_block.var(var_name)
|
||||
var_shape = var.shape
|
||||
var_dtype = var.dtype
|
||||
assert str(var_dtype) in supported_var_dtypes, (
|
||||
'Found unsupported variable dtype: "{}", current supported '
|
||||
'dtype(s) is/are: [{}]. '.format(
|
||||
str(var_dtype), ", ".join(supported_var_dtypes)
|
||||
)
|
||||
)
|
||||
(
|
||||
logger.info(
|
||||
f'[+] var: "{var_name}", shape={var_shape}, dtype="{var_dtype}".\n'
|
||||
)
|
||||
if verbose
|
||||
else None
|
||||
)
|
||||
np_dtype = (
|
||||
convert_dtype(var_dtype)
|
||||
if isinstance(var_dtype, core.VarDesc.VarType)
|
||||
else var_dtype
|
||||
)
|
||||
if str(var_dtype).find('int') != -1:
|
||||
# target variable's type is int* (uint*, int*), it is highly possible that
|
||||
# the target variable contains indices (such as lookup_table op's input var)
|
||||
# for safety we need to fill it with all one instead of random numbers
|
||||
# NOTE (liuchenghao): filling with zero will generate "division by zero" error
|
||||
# in mod ops, so filling with one seems to be the simplest way to make it work,
|
||||
# although it is possible that for array with only one element, index "1" is
|
||||
# invalid, that situation is very rare and we don't need to care about it now.
|
||||
new_tensor = np.array(np.ones(var_shape)).astype(np_dtype)
|
||||
else:
|
||||
# target variable's type is float*, we treat it as an ordinary tensor, fill it
|
||||
# with random gaussian numbers
|
||||
new_tensor = np.array(np.random.randn(*var_shape)).astype(np_dtype)
|
||||
new_tensor = _as_lodtensor(new_tensor, place, var_dtype)
|
||||
check_feed_shape_type(var, new_tensor)
|
||||
core.set_variable(scope, new_tensor, var_name)
|
||||
return new_tensor
|
||||
|
||||
def _configure_feed_ops_and_return_feed_names():
|
||||
"""
|
||||
configure feed op,
|
||||
1. alloc feed op output var storage
|
||||
2. fill feed op's input var
|
||||
return feed var names
|
||||
"""
|
||||
|
||||
feed_names = []
|
||||
has_feed_op = False
|
||||
for op in cloned_main_block.ops:
|
||||
if op.type == "feed":
|
||||
has_feed_op = True
|
||||
out_var_name = op.desc.output('Out')[0]
|
||||
in_var_name = op.desc.input('X')[0] # this is usually "feed"
|
||||
input_index = op.desc.attr('col')
|
||||
new_tensor = _alloc_and_fill_var(out_var_name)
|
||||
core.set_feed_variable(
|
||||
scope, new_tensor, in_var_name, input_index
|
||||
)
|
||||
feed_names.append(out_var_name)
|
||||
if not has_feed_op:
|
||||
(
|
||||
logger.info("WARNING: program does not have any feed op.\n")
|
||||
if verbose
|
||||
else None
|
||||
)
|
||||
return feed_names
|
||||
|
||||
for var_name in _analyze_graph_and_collect_all_vars_with_zero_in_degree():
|
||||
_alloc_and_fill_var(var_name)
|
||||
feed_names = _configure_feed_ops_and_return_feed_names()
|
||||
|
||||
# build a simple plan from program and run profiling
|
||||
plan = core.Plan([core.Job("default")], {"default": cloned_program.desc})
|
||||
exe = _StandaloneExecutor(place, plan, scope)
|
||||
|
||||
num_ops = len(cloned_main_block.ops)
|
||||
prof_results = [[None for _ in range(run_iters)] for _ in range(num_ops)]
|
||||
|
||||
for iter_id in range(run_iters):
|
||||
# for each iteration, run profiling and retrieve modified version of program desc
|
||||
program_desc = exe.run_profile(feed_names)
|
||||
|
||||
# rebuild program object from the new program desc
|
||||
temp_program = cloned_program.clone()
|
||||
temp_program._rebuild_from_desc(program_desc)
|
||||
temp_main_block = temp_program.global_block()
|
||||
|
||||
# collect profiling result
|
||||
for op_id, temp_op in zip(
|
||||
range(len(temp_main_block.ops)), temp_main_block.ops
|
||||
):
|
||||
run_time_us = temp_op.dist_attr.run_time_us
|
||||
prof_results[op_id][iter_id] = (
|
||||
run_time_us
|
||||
if check_if_op_supports_runtime_profiling(temp_op)
|
||||
and run_time_us >= 0.0
|
||||
else None
|
||||
)
|
||||
return prof_results
|
||||
|
||||
|
||||
def measure_program_real_op_cost(
|
||||
program: paddle.static.Program,
|
||||
run_iters: int = 8,
|
||||
place=paddle.base.framework._current_expected_place(),
|
||||
verbose_level: int = 0,
|
||||
):
|
||||
'''
|
||||
Description
|
||||
-----------
|
||||
Measuring real op run time (us) with respect to the given "program" and "place".
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
@param program: paddle.static.Program
|
||||
The program object waiting to be executed.
|
||||
@param run_iters: int
|
||||
Specify how many iterations will be run during profiling. Larger value tends
|
||||
to give more accurate profiling result but requires more time.
|
||||
@param place: paddle.CPUPlace | paddle.CUDAPlace
|
||||
Where the program is going to be executed.
|
||||
@param verbose_level: int
|
||||
Set up verbose level during profiling. Can be set to one of the following:
|
||||
0 = turn off, don't output anything,
|
||||
1 = output profiling messages only,
|
||||
2 = output profiling and debug messages.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
Nothing to return. This API will write op run time directly into program object.
|
||||
For example, to retrieve the run time for the first op in program, use:
|
||||
>>> program.global_block().ops[0].dist_attr.run_time_us
|
||||
|
||||
Note
|
||||
-----------
|
||||
Not all ops support runtime profiling. Currently communication ops do not support
|
||||
runtime profiling feature since their execution times rely on other ops. To check
|
||||
if an op supports runtime profiling, use:
|
||||
>>> check_if_op_supports_runtime_profiling(op)
|
||||
where "op" is an instance of "paddle.base.framework.Operator".
|
||||
|
||||
Example
|
||||
-----------
|
||||
* Profiling a simple program from scratch:
|
||||
>>> from paddle.distributed.auto_parallel.static.utils import (
|
||||
... measure_program_real_op_cost,
|
||||
... )
|
||||
>>> program = ... # build your own program object here.
|
||||
>>> measure_program_real_op_cost(
|
||||
>>> program, verbose_level=1
|
||||
>>> )
|
||||
* Profiling a program which is already embedded into an Executor or some other class instance:
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed.auto_parallel.static.utils import (
|
||||
... measure_program_real_op_cost,
|
||||
... )
|
||||
>>> place: str = paddle.device.get_device() # here we assume place = "cuda:x"
|
||||
>>> place = paddle.CUDAPlace(int(place.split(':')[1]))
|
||||
>>> # here "program" is an inner object that has already been built before
|
||||
>>> measure_program_real_op_cost(program, verbose_level=1)
|
||||
'''
|
||||
|
||||
assert isinstance(program, Program), (
|
||||
f'"program" should be a instance of "paddle.base.framework.Program" but got type "{type(program).__name__}".'
|
||||
)
|
||||
supported_places = [
|
||||
paddle.CUDAPlace,
|
||||
]
|
||||
assert any(
|
||||
isinstance(place, supported_place)
|
||||
for supported_place in supported_places
|
||||
), (
|
||||
f'Current place ({place}) does not support runtime profiling. "place" should be one of the following: {supported_places}.'
|
||||
)
|
||||
assert isinstance(run_iters, int) and run_iters >= 1, (
|
||||
'Invalid parameter run_iters set. run_iters should be an integer >= 1.'
|
||||
)
|
||||
if run_iters == 1:
|
||||
warnings.warn(
|
||||
'run_iters was set to 1, profiling results might be inaccurate due to outliers.'
|
||||
)
|
||||
|
||||
logger = get_logger(log_level=logging.INFO)
|
||||
|
||||
# run profiling multiple times and record op run time of each run
|
||||
prof_results = _measure_program_real_op_cost_multipass(
|
||||
program, place, run_iters, verbose=(verbose_level >= 2)
|
||||
)
|
||||
op_num = len(prof_results)
|
||||
for op_id, op in zip(range(op_num), program.global_block().ops):
|
||||
op_runtime_us_final = None
|
||||
if prof_results[op_id][0] is not None:
|
||||
op_runtime_us_final = np.median(prof_results[op_id])
|
||||
if (
|
||||
op_runtime_us_final is not None
|
||||
and check_if_op_supports_runtime_profiling(op)
|
||||
):
|
||||
op.dist_attr.run_time_us = op_runtime_us_final
|
||||
(
|
||||
logger.info(
|
||||
f"{op_id!s:>4} {op.type!s:>32} {op_runtime_us_final:.1f} us"
|
||||
)
|
||||
if verbose_level >= 1
|
||||
else None
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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 paddle
|
||||
from paddle.distributed.auto_parallel.static.dist_tensor import (
|
||||
DistributedTensor,
|
||||
)
|
||||
from paddle.static import Variable
|
||||
|
||||
from .base_cost import Cost
|
||||
|
||||
|
||||
class TensorCost:
|
||||
def __init__(self, tensor=None, dist_tensor=None, shape=None, dtype=None):
|
||||
self._check_args(tensor, dist_tensor, shape, dtype)
|
||||
self._tensor = tensor
|
||||
self._dist_tensor = dist_tensor
|
||||
self._shape = shape
|
||||
self._dtype = dtype
|
||||
self._cost = self.calc_cost()
|
||||
|
||||
@property
|
||||
def tensor(self):
|
||||
return self._tensor
|
||||
|
||||
@property
|
||||
def dist_tensor(self):
|
||||
return self._dist_tensor
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._shape
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._dtype
|
||||
|
||||
def _check_args(self, tensor, dist_tensor, shape, dtype):
|
||||
if tensor is not None:
|
||||
assert shape is None and dist_tensor is None and dtype is None
|
||||
|
||||
if not isinstance(tensor, Variable):
|
||||
raise TypeError(
|
||||
f"Please check tensor type is Variable, but got {type(tensor)}"
|
||||
)
|
||||
|
||||
elif dist_tensor is not None:
|
||||
assert tensor is None and shape is None
|
||||
if not isinstance(dist_tensor, DistributedTensor):
|
||||
raise TypeError(
|
||||
f"Please check dist_tensor type is DistributedTensor, but got {type(dist_tensor)}"
|
||||
)
|
||||
|
||||
elif shape is not None:
|
||||
assert tensor is None and dist_tensor is None and dtype is not None
|
||||
if not isinstance(shape, (list, set)):
|
||||
raise TypeError(
|
||||
f"Please check shape type is list or set, but got {type(shape)}"
|
||||
)
|
||||
|
||||
elif dtype is not None:
|
||||
assert tensor is None and dist_tensor is None and shape is not None
|
||||
|
||||
@property
|
||||
def cost(self):
|
||||
return self._cost
|
||||
|
||||
def calc_cost(self):
|
||||
dtype = None
|
||||
shape = None
|
||||
|
||||
if self.dist_tensor:
|
||||
shape = self.dist_tensor.local_sizes()
|
||||
dtype = self.dist_tensor.serial_tensor.dtype
|
||||
elif self.tensor:
|
||||
shape = self.tensor.shape
|
||||
dtype = self.tensor.dtype
|
||||
elif self.shape and self.dtype:
|
||||
shape = self.shape
|
||||
dtype = self.dtype
|
||||
|
||||
total_count = reduce(lambda x, y: x * y, shape, 1)
|
||||
|
||||
if dtype == paddle.float32 or dtype == paddle.int32:
|
||||
dtype_factor = 4
|
||||
elif dtype == paddle.int64:
|
||||
dtype_factor = 8
|
||||
elif dtype == paddle.uint8:
|
||||
dtype_factor = 1
|
||||
else:
|
||||
dtype_factor = 2
|
||||
|
||||
memory = total_count * dtype_factor
|
||||
assert memory >= 0
|
||||
cost = Cost(memory=memory)
|
||||
|
||||
return cost
|
||||
Reference in New Issue
Block a user