wehub-resource-sync
544 lines
20 KiB
Python
544 lines
20 KiB
Python
# 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 copy
|
|
from collections import defaultdict
|
|
|
|
import paddle
|
|
from paddle.distributed.auto_parallel.static.dist_context import (
|
|
DistributedContext,
|
|
)
|
|
from paddle.distributed.auto_parallel.static.operators.common import (
|
|
get_distributed_operator_impl_container,
|
|
)
|
|
from paddle.framework import Program, core
|
|
from paddle.static import Parameter
|
|
|
|
from .dist_attribute import OperatorDistAttr
|
|
from .operators.common import BACKWARD_ONLY_DIST_OPS
|
|
from .utils import (
|
|
__no_shape_var_type__,
|
|
is_backward_op,
|
|
is_forward_op,
|
|
is_loss_op,
|
|
is_optimize_op,
|
|
)
|
|
|
|
__varname_not_in_block__ = ["lod_tensor_blocking_queue"]
|
|
|
|
|
|
class Partitioner:
|
|
"""
|
|
warning:: Partitioner is experimental and subject to change.
|
|
|
|
Partitioner convert a program into another program.
|
|
Given a serial program which has been auto completed with shard annotation, the Partitioner
|
|
convert the serial program into a "distributed" program. The Partitioner will modify the serial
|
|
program in following two ways, which is also the major difference between serial and distributed program:
|
|
1. partition op: replace a serial op into its corresponding dist op inferred from the shard annotation
|
|
2. partition var: if a var is sharded, modify the shape of var according to its shard annotation
|
|
|
|
Partitioner is supposed to be call by the auto parallel framework, and not supposed to be directly called by user.
|
|
"""
|
|
|
|
def __init__(self, dist_context, rank_id=0):
|
|
"""
|
|
Args:
|
|
dist_context (DistributedContext): used to access the distributed_attr of var & op, every Partitioner object could maintain its own DistributedContext member, and partition program base on that shard scenario.
|
|
rank_id (int): global rank id to which the partitioned distributed program belong.
|
|
"""
|
|
if not isinstance(dist_context, DistributedContext):
|
|
raise TypeError(
|
|
f"dist_context be DistributedContext, got {type(dist_context)} here"
|
|
)
|
|
|
|
self._dist_context = dist_context
|
|
self._rank_id = rank_id
|
|
self._serial2dist_varname_mapping = defaultdict(
|
|
dict
|
|
) # block_id -> serial_varname -> dist_varname
|
|
self._dist_varname_suffix = ""
|
|
self._forward_op_id2forward_op = {}
|
|
|
|
def partition(
|
|
self, serial_main_program, serial_startup_program, params_grads
|
|
):
|
|
if not isinstance(serial_main_program, (Program)):
|
|
raise TypeError(
|
|
f"main_program be paddle.framework.Program, got {type(serial_main_program)} here"
|
|
)
|
|
|
|
# check if shard annotated serial program valid
|
|
if not self._is_valid_annotated_program(serial_main_program):
|
|
raise RuntimeError(
|
|
"Not all vars or ops are annotated in main program !"
|
|
)
|
|
|
|
# init distop helper
|
|
dist_op_context = self._dist_context.dist_op_context
|
|
dist_op_context.varname_mapping = self._serial2dist_varname_mapping
|
|
dist_op_context.rank_id = self._rank_id
|
|
|
|
# partition startup program
|
|
if serial_startup_program is None:
|
|
partitioned_startup_prog = None
|
|
else:
|
|
partitioned_startup_prog = self.partition_startup_program(
|
|
serial_main_program, serial_startup_program
|
|
)
|
|
dist_op_context.dst_startup_program = partitioned_startup_prog
|
|
|
|
# partition main program
|
|
(
|
|
partitioned_main_prog,
|
|
partitioned_params_grads,
|
|
) = self.partition_main_program(serial_main_program, params_grads)
|
|
|
|
return (
|
|
partitioned_main_prog,
|
|
partitioned_startup_prog,
|
|
partitioned_params_grads,
|
|
)
|
|
|
|
def partition_startup_program(
|
|
self, serial_main_program, serial_startup_program
|
|
):
|
|
if not isinstance(serial_startup_program, (Program)):
|
|
raise TypeError(
|
|
f"dist_context be paddle.framework.Program, got {type(serial_startup_program)} here"
|
|
)
|
|
|
|
partitioned_startup_prog = paddle.framework.Program()
|
|
partitioned_startup_prog._name_generator = (
|
|
serial_startup_program._name_generator.clone()
|
|
)
|
|
ref_block = serial_main_program.global_block()
|
|
target_block = partitioned_startup_prog.global_block()
|
|
var2shape = {}
|
|
temp_varname_map = {}
|
|
|
|
# tensors
|
|
for var in serial_startup_program.list_vars():
|
|
assert var.persistable
|
|
new_name = var.name + self._dist_varname_suffix
|
|
temp_varname_map[var.name] = new_name
|
|
target_shape = _partition_var(
|
|
self._dist_context, ref_block, target_block, var.name, new_name
|
|
)
|
|
var2shape[new_name] = target_shape
|
|
|
|
# ops
|
|
for op in serial_startup_program.global_block().ops:
|
|
# TODO if var not belong to this rank, should be filtered
|
|
output_vars = op.desc.output_arg_names()
|
|
assert len(output_vars) == 1, (
|
|
f"initializer should output only ONE variable, but got [{op.desc}]"
|
|
)
|
|
assert temp_varname_map[output_vars[0]] in var2shape, (
|
|
f"try to initialize [{output_vars[0]}] which is not a persistable var"
|
|
)
|
|
new_op_desc = target_block.desc.append_op()
|
|
new_op_desc.copy_from(op.desc)
|
|
new_op_desc._rename_output(
|
|
output_vars[0], temp_varname_map[output_vars[0]]
|
|
)
|
|
new_op_desc._set_attr(
|
|
"shape", var2shape[temp_varname_map[output_vars[0]]]
|
|
)
|
|
target_block._sync_with_cpp()
|
|
|
|
# set distribute attribute
|
|
new_op = target_block.ops[-1]
|
|
assert new_op.type == new_op_desc.type()
|
|
assert new_op.desc == new_op_desc
|
|
output_var = target_block.var(output_vars[0])
|
|
output_var_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(output_var)
|
|
)
|
|
op_attr = OperatorDistAttr()
|
|
op_attr.process_mesh = output_var_attr.process_mesh
|
|
op_attr.set_output_dims_mapping(
|
|
output_var.name, output_var_attr.dims_mapping
|
|
)
|
|
op_attr.set_input_dims_mapping(
|
|
output_var.name, output_var_attr.dims_mapping
|
|
)
|
|
self._dist_context.set_op_dist_attr_for_program(new_op, op_attr)
|
|
|
|
return partitioned_startup_prog
|
|
|
|
def partition_main_program(self, serial_main_program, params_and_grads):
|
|
"""
|
|
1. partition variables
|
|
2. replace local op with corresponding dist op
|
|
"""
|
|
|
|
partitioned_main_prog = paddle.framework.Program()
|
|
partitioned_main_prog._name_generator = (
|
|
serial_main_program._name_generator.clone()
|
|
)
|
|
dist_op_context = self._dist_context.dist_op_context
|
|
dist_op_context.dst_main_program = partitioned_main_prog
|
|
|
|
for idx in range(self._dist_context.block_state.nblock):
|
|
ref_block = serial_main_program.blocks[idx]
|
|
|
|
if idx == 0:
|
|
target_block = partitioned_main_prog.blocks[0]
|
|
else:
|
|
target_block = partitioned_main_prog._create_block(
|
|
parent_idx=ref_block.parent_idx
|
|
)
|
|
assert ref_block.idx == target_block.idx
|
|
target_block._set_forward_block_idx(ref_block.forward_block_idx)
|
|
dist_op_context.work_block = target_block
|
|
self.partition_block(ref_block, target_block)
|
|
|
|
partitioned_main_prog.current_block_idx = 0
|
|
|
|
# should reconnect the block_attr ptr to the correct block
|
|
for block_id in range(self._dist_context.block_state.nblock):
|
|
block = partitioned_main_prog.block(block_id)
|
|
for op in block.ops:
|
|
for attr_name in op.all_attrs():
|
|
if op.attr_type(attr_name) == core.AttrType.BLOCK:
|
|
relative_id = op._block_attr_id(attr_name)
|
|
op._set_attr(
|
|
attr_name, partitioned_main_prog.block(relative_id)
|
|
)
|
|
|
|
partitioned_params_and_grads = []
|
|
for p, g in params_and_grads:
|
|
assert p.name in self._serial2dist_varname_mapping[0]
|
|
dist_p = self._get_dist_var_by_serial_var(
|
|
p, partitioned_main_prog, 0
|
|
)
|
|
if g is None:
|
|
dist_g = None
|
|
else:
|
|
assert g.name in self._serial2dist_varname_mapping[0]
|
|
dist_g = self._get_dist_var_by_serial_var(
|
|
g, partitioned_main_prog, 0
|
|
)
|
|
partitioned_params_and_grads.append((dist_p, dist_g))
|
|
|
|
return partitioned_main_prog, partitioned_params_and_grads
|
|
|
|
def partition_block(self, ref_block, target_block):
|
|
dist_op_context = self._dist_context.dist_op_context
|
|
|
|
last_fwd_op_idx = -1
|
|
for idx, op in enumerate(ref_block.ops):
|
|
if is_loss_op(op):
|
|
last_fwd_op_idx = idx
|
|
break
|
|
|
|
if last_fwd_op_idx == -1:
|
|
last_fwd_op_idx = len(ref_block.ops)
|
|
|
|
for idx in range(len(ref_block.ops)):
|
|
if idx <= last_fwd_op_idx:
|
|
self._forward_op_id2forward_op[
|
|
ref_block.ops[idx].desc.original_id()
|
|
] = ref_block.ops[idx]
|
|
|
|
# partition
|
|
appended_grad_times = 0
|
|
for idx, op in enumerate(ref_block.ops):
|
|
op_dist_attr = self._dist_context.get_op_dist_attr_for_program(op)
|
|
if is_backward_op(op) and (
|
|
is_forward_op(ref_block.ops[idx - 1])
|
|
or is_loss_op(ref_block.ops[idx - 1])
|
|
):
|
|
if not op_dist_attr.is_recompute:
|
|
appended_grad_times += 1
|
|
|
|
# partition input variables
|
|
for serial_input_varname in op.desc.input_arg_names():
|
|
if (
|
|
serial_input_varname
|
|
not in self._serial2dist_varname_mapping[
|
|
ref_block.forward_block_idx
|
|
]
|
|
or serial_input_varname
|
|
not in self._serial2dist_varname_mapping[ref_block.idx]
|
|
):
|
|
new_varname = (
|
|
serial_input_varname + self._dist_varname_suffix
|
|
)
|
|
if ref_block.has_var(serial_input_varname):
|
|
_partition_var(
|
|
self._dist_context,
|
|
ref_block,
|
|
target_block,
|
|
serial_input_varname,
|
|
new_varname,
|
|
)
|
|
self._serial2dist_varname_mapping[ref_block.idx][
|
|
serial_input_varname
|
|
] = new_varname
|
|
|
|
# partition output vars
|
|
for serial_output_varname in op.desc.output_arg_names():
|
|
if (
|
|
serial_output_varname
|
|
not in self._serial2dist_varname_mapping[
|
|
ref_block.forward_block_idx
|
|
]
|
|
or serial_output_varname
|
|
not in self._serial2dist_varname_mapping[ref_block.idx]
|
|
):
|
|
new_varname = (
|
|
serial_output_varname + self._dist_varname_suffix
|
|
)
|
|
if ref_block.has_var(serial_output_varname):
|
|
_partition_var(
|
|
self._dist_context,
|
|
ref_block,
|
|
target_block,
|
|
serial_output_varname,
|
|
new_varname,
|
|
)
|
|
self._serial2dist_varname_mapping[ref_block.idx][
|
|
serial_output_varname
|
|
] = new_varname
|
|
|
|
# partition op
|
|
if is_forward_op(op) or op_dist_attr.is_recompute:
|
|
kinputs, koutputs = dist_op_context.prepare_context(op)
|
|
dist_op_forward_impl = _get_dist_op_forward_implement(
|
|
op, self._dist_context
|
|
)
|
|
dist_op_forward_impl.forward(
|
|
self._dist_context, **kinputs, **koutputs
|
|
)
|
|
|
|
elif is_backward_op(op):
|
|
kinputs, koutputs = dist_op_context.prepare_context(op)
|
|
dist_op_backward_impl = _get_dist_op_backward_implement(
|
|
op, self._dist_context, self._forward_op_id2forward_op
|
|
)
|
|
grad_var_to_var = (
|
|
self._dist_context.dist_op_context.grad_var_to_var[
|
|
appended_grad_times
|
|
]
|
|
)
|
|
dist_op_backward_impl.backward(
|
|
self._dist_context,
|
|
**kinputs,
|
|
**koutputs,
|
|
**{"grad_var_to_var": grad_var_to_var},
|
|
)
|
|
elif is_optimize_op(op):
|
|
# NOTE: BACKWARD_ONLY_DIST_OPS's op_role must be 2 because of 1F1B PASS
|
|
kinputs, koutputs = dist_op_context.prepare_context(op)
|
|
dist_op_opt_impl = _get_dist_op_backward_implement(
|
|
op, self._dist_context, self._forward_op_id2forward_op
|
|
)
|
|
dist_op_opt_impl.backward(
|
|
self._dist_context,
|
|
**kinputs,
|
|
**koutputs,
|
|
**{"grad_var_to_var": {}},
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"partitioner only support forward and backward, optimize ops, but got {op}"
|
|
)
|
|
|
|
def _is_valid_annotated_program(self, program):
|
|
# TODO (ZJ-LIANG) should check all block
|
|
ops = program.global_block().ops
|
|
vars_ = program.list_vars()
|
|
op_dist_attrs = [
|
|
self._dist_context.get_op_dist_attr_for_program(op) for op in ops
|
|
]
|
|
var_dist_attrs = [
|
|
self._dist_context.get_tensor_dist_attr_for_program(var)
|
|
for var in vars_
|
|
if (var.type not in __no_shape_var_type__)
|
|
]
|
|
|
|
all_ops_annotated = all(
|
|
dist_attr is not None for dist_attr in op_dist_attrs
|
|
)
|
|
all_vars_annotated = all(
|
|
dist_attr is not None for dist_attr in var_dist_attrs
|
|
)
|
|
|
|
return all_ops_annotated and all_vars_annotated
|
|
|
|
def _get_dist_var_by_serial_var(
|
|
self, serial_var, partitioned_main_prog, block_id
|
|
):
|
|
block_idx = serial_var.block.idx
|
|
target_block = partitioned_main_prog.blocks[block_idx]
|
|
dist_var_name = self._serial2dist_varname_mapping[block_id][
|
|
serial_var.name
|
|
]
|
|
assert target_block.has_var(dist_var_name)
|
|
return target_block.var(dist_var_name)
|
|
|
|
|
|
def _get_dist_shape(var, dist_attr):
|
|
var_shape = var.shape
|
|
mapping = dist_attr.dims_mapping
|
|
mesh = dist_attr.process_mesh.shape
|
|
if mapping == []:
|
|
return var_shape
|
|
|
|
assert len(var_shape) == len(mapping), (
|
|
f"variable shape [{var_shape}] and dim_mapping [{mapping}] is NOT match !"
|
|
)
|
|
new_shape = []
|
|
for idx in range(len(var_shape)):
|
|
if var_shape[idx] == -1 or mapping[idx] == -1:
|
|
new_shape.append(var_shape[idx])
|
|
else:
|
|
assert var_shape[idx] % mesh[mapping[idx]] == 0, (
|
|
f"un-event partition: var_shape[idx]=[{var_shape[idx]}], mesh[{mesh[mapping[idx]]}], {var.name}, {var_shape}, {mesh}, {mapping}"
|
|
)
|
|
new_shape.append(var_shape[idx] // mesh[mapping[idx]])
|
|
|
|
return new_shape
|
|
|
|
|
|
def _partition_parameter(
|
|
dist_context, src_var, dst_block, dst_varname, dst_shape
|
|
):
|
|
# NOTE hack to copied Parameter
|
|
# not initialized parameter, need to initialize it
|
|
copied_kwargs = {}
|
|
copied_kwargs['trainable'] = src_var.trainable
|
|
copied_kwargs['optimize_attr'] = src_var.optimize_attr
|
|
copied_kwargs['regularizer'] = src_var.regularizer
|
|
copied_kwargs['do_model_average'] = src_var.do_model_average
|
|
copied_kwargs['need_clip'] = src_var.need_clip
|
|
|
|
param = Parameter(
|
|
block=dst_block,
|
|
type=src_var.type,
|
|
name=dst_varname,
|
|
shape=dst_shape,
|
|
dtype=src_var.dtype,
|
|
lod_level=src_var.lod_level,
|
|
error_clip=src_var.error_clip,
|
|
stop_gradient=src_var.stop_gradient,
|
|
is_data=src_var.is_data,
|
|
belong_to_optimizer=src_var.belong_to_optimizer,
|
|
**copied_kwargs,
|
|
)
|
|
|
|
return param
|
|
|
|
|
|
def _partition_intermediate_var(
|
|
dist_context, src_var, dst_block, dst_varname, dst_shape
|
|
):
|
|
var = dst_block.create_var(
|
|
type=src_var.type,
|
|
name=dst_varname,
|
|
shape=dst_shape,
|
|
dtype=src_var.dtype,
|
|
lod_level=src_var.lod_level,
|
|
persistable=src_var.persistable,
|
|
error_clip=src_var.error_clip,
|
|
stop_gradient=src_var.stop_gradient,
|
|
is_data=src_var.is_data,
|
|
belong_to_optimizer=src_var.belong_to_optimizer,
|
|
)
|
|
|
|
return var
|
|
|
|
|
|
def _partition_var(
|
|
dist_context, src_block, dst_block, src_varname, dst_varname
|
|
):
|
|
"""
|
|
partition include: split + replicate
|
|
"""
|
|
src_var = src_block.var(src_varname)
|
|
|
|
if src_var.type in __no_shape_var_type__:
|
|
persist = getattr(src_var, 'persistable', False)
|
|
new_var = dst_block.create_var(
|
|
type=src_var.type,
|
|
name=dst_varname,
|
|
persistable=persist,
|
|
stop_gradient=True,
|
|
)
|
|
target_shape = None
|
|
else:
|
|
dist_attr = dist_context.get_tensor_dist_attr_for_program(src_var)
|
|
target_shape = _get_dist_shape(src_var, dist_attr)
|
|
|
|
if isinstance(src_var, Parameter):
|
|
new_var = _partition_parameter(
|
|
dist_context, src_var, dst_block, dst_varname, target_shape
|
|
)
|
|
else:
|
|
new_var = _partition_intermediate_var(
|
|
dist_context, src_var, dst_block, dst_varname, target_shape
|
|
)
|
|
|
|
dist_attr = copy.deepcopy(
|
|
dist_context.get_tensor_dist_attr_for_program(src_var)
|
|
)
|
|
assert dist_attr is not None
|
|
dist_context.set_tensor_dist_attr_for_program(new_var, dist_attr)
|
|
|
|
return target_shape
|
|
|
|
|
|
def _get_dist_op_backward_implement(
|
|
backward_op, dist_context, forward_op_id2forward_op
|
|
):
|
|
dist_op_context = dist_context.dist_op_context
|
|
if backward_op.desc.original_id() in dist_op_context.grad_op_id_to_op_id:
|
|
forward_op_id = dist_op_context.grad_op_id_to_op_id[
|
|
backward_op.desc.original_id()
|
|
]
|
|
forward_op = forward_op_id2forward_op[forward_op_id]
|
|
forward_op_dist_attr = dist_context.get_op_dist_attr_for_program(
|
|
forward_op
|
|
)
|
|
dist_op_impl_container = get_distributed_operator_impl_container(
|
|
forward_op_dist_attr.impl_type
|
|
)
|
|
dist_op_impl = dist_op_impl_container.get_impl(
|
|
forward_op_dist_attr.impl_idx
|
|
)
|
|
return dist_op_impl
|
|
|
|
# # NOTE trick for dist ops that only have backward implement
|
|
if backward_op.type in BACKWARD_ONLY_DIST_OPS:
|
|
op_dist_attr = dist_context.get_op_dist_attr_for_program(backward_op)
|
|
assert op_dist_attr.impl_idx >= 0
|
|
dist_op_impl = get_distributed_operator_impl_container(
|
|
op_dist_attr.impl_type
|
|
).get_impl(op_dist_attr.impl_idx)
|
|
return dist_op_impl
|
|
|
|
dist_op = get_distributed_operator_impl_container("default")
|
|
return dist_op.get_impl(0)
|
|
|
|
|
|
def _get_dist_op_forward_implement(forward_op, dist_context):
|
|
dist_attr = dist_context.get_op_dist_attr_for_program(forward_op)
|
|
dist_op_impl_container = get_distributed_operator_impl_container(
|
|
dist_attr.impl_type
|
|
)
|
|
dist_op_impl = dist_op_impl_container.get_impl(dist_attr.impl_idx)
|
|
return dist_op_impl
|