Skip to content

P
Paddle

BaiXuePrincess / Paddle
与 Fork 源项目一致

Fork自 PaddlePaddle / Paddle

通知 1
Star 1
Fork 0
P
Paddle
You need to sign in or sign up before continuing.
未验证 提交 3a014783 编写于 作者: Y Yuang Liu 提交者： GitHub
浏览文件
操作

update the split logic for uniform (#47670) (#47705) 

* code format change

* update the split logic for uniform (#47670)
上级 d5809836
显示空白变更内容
内联 并排
Showing with 221 addition and 130 deletion
python/paddle/distributed/fleet/meta_parallel/parallel_layers/pp_layers.py
浏览文件 @ 3a014783
......@@ -57,7 +57,6 @@ __all__ = []
class LayerDesc(object):
def __init__(self, layer_func, *inputs, **kwargs):
self.layer_func = layer_func
self.inputs = inputs
......@@ -65,25 +64,28 @@ class LayerDesc(object):
if not issubclass(layer_func, Layer):
raise TypeError(
"The input(layer_func) should be a derived class of Layer.")
"The input(layer_func) should be a derived class of Layer."
)
def build_layer(self):
return self.layer_func(*self.inputs, **self.kwargs)
def __repr__(self):
return layer_to_str(self.layer_func.__name__, *self.inputs,
**self.kwargs)
return layer_to_str(
self.layer_func.__name__, *self.inputs, **self.kwargs
)
class SharedLayerDesc(LayerDesc):
def __init__(self,
def __init__(
self,
key,
layer_func,
forward_func=None,
shared_weight_attr='weight',
*inputs,
**kwargs):
**kwargs
):
super(SharedLayerDesc, self).__init__(layer_func, *inputs, **kwargs)
self.layer_name = key
self.forward_func = forward_func
......@@ -91,12 +93,13 @@ class SharedLayerDesc(LayerDesc):
class SegmentLayers(object):
def __init__(self,
def __init__(
self,
layers_desc,
num_parts,
method="uniform",
num_virtual_pipeline_stage=None):
num_virtual_pipeline_stage=None,
):
self._layers_desc = layers_desc
self.method = method
self.num_parts = num_parts
......@@ -104,7 +107,9 @@ class SegmentLayers(object):
self.num_virtual_pipeline_stage = num_virtual_pipeline_stage
if self.num_virtual_pipeline_stage is not None:
self.total_parts = num_parts * self.num_virtual_pipeline_stage
assert self.num_items >= self.num_parts, "layer number should be greater than number of segments"
assert (
self.num_items >= self.num_parts
), "layer number should be greater than number of segments"
def do_segment(self):
if self.method == "uniform":
......@@ -118,12 +123,17 @@ class SegmentLayers(object):
for idx in weight_idxs:
weights[idx] = 1
actual_num_parts = self.num_parts if self.num_virtual_pipeline_stage is None else self.total_parts
actual_num_parts = (
self.num_parts
if self.num_virtual_pipeline_stage is None
else self.total_parts
)
assert sum(
weights
) % actual_num_parts == 0, "number of layers ({}) should be divided by part number({})".format(
sum(weights), actual_num_parts)
assert (
sum(weights) % actual_num_parts == 0
), "number of layers ({}) should be divided by part number({})".format(
sum(weights), actual_num_parts
)
part_size = sum(weights) // actual_num_parts
result = [0 for _ in range(actual_num_parts + 1)]
......@@ -156,21 +166,23 @@ class SegmentLayers(object):
if regex.search(name):
weight_idxs.append(idx)
assert len(
weight_idxs) > 0, "weight_idxs' length should be greater than 0"
assert (
len(weight_idxs) > 0
), "weight_idxs' length should be greater than 0"
return weight_idxs
def uniform(self, num_items, num_parts):
result = [0 for _ in range(num_parts + 1)]
part_size = math.floor(num_items / num_parts)
for i in range(num_parts):
result[i] = int(min(part_size * i, num_items))
extra_layers = num_items % num_parts
for i in range(1, num_parts):
offset = 1 if i > (num_parts - extra_layers) else 0
result[i] = int(min(result[i - 1] + part_size + offset, num_items))
result[num_parts] = num_items
return result
class PipelineLayerChunk(Layer):
def __init__(self):
super(PipelineLayerChunk, self).__init__()
self.run_function = []
......@@ -192,7 +204,8 @@ class PipelineLayerChunk(Layer):
# behavior under recompute circumstance.
raise PermissionError(
"The forward function of PipelineLayerChunk cannot be called directly. "
"Please call forward function of PipelineLayer.")
"Please call forward function of PipelineLayer."
)
class PipelineLayer(Layer):
......@@ -274,7 +287,8 @@ class PipelineLayer(Layer):
"""
def __init__(self,
def __init__(
self,
layers,
num_stages=None,
topology=None,
......@@ -282,24 +296,32 @@ class PipelineLayer(Layer):
seg_method="uniform",
recompute_interval=0,
recompute_ctx=None,
num_virtual_pipeline_stages=None):
num_virtual_pipeline_stages=None,
):
super(PipelineLayer, self).__init__()
if num_stages is None and topology is None:
raise ValueError("should provide num_stages or topology")
if num_virtual_pipeline_stages:
assert isinstance(num_virtual_pipeline_stages, int), \
"virtual_pipeline_stage should be None or an int"
assert isinstance(
num_virtual_pipeline_stages, int
), "virtual_pipeline_stage should be None or an int"
if num_virtual_pipeline_stages > 1:
logger.info(
"set num_virtual_pipeline_stages > 1 means using interleave scheduler instead of 1f1b scheduler"
)
assert isinstance(seg_method, str), \
"seg_method should be a str for interleave scheduler"
assert seg_method.startswith('layer:'), \
"seg_method shoud be start with layer: for interleave scheduler"
self._num_virtual_pipeline_stages = 1 if num_virtual_pipeline_stages is None else num_virtual_pipeline_stages
assert isinstance(
seg_method, str
), "seg_method should be a str for interleave scheduler"
assert seg_method.startswith(
'layer:'
), "seg_method shoud be start with layer: for interleave scheduler"
self._num_virtual_pipeline_stages = (
1
if num_virtual_pipeline_stages is None
else num_virtual_pipeline_stages
)
# lazy import
import paddle.distributed as dist
......@@ -313,13 +335,17 @@ class PipelineLayer(Layer):
self.recompute_ctx = recompute_ctx
if recompute_interval > 0:
assert recompute_ctx is not None, "recompute_ctx must be not None for recompute."
assert (
recompute_ctx is not None
), "recompute_ctx must be not None for recompute."
offload = recompute_ctx.get('offload', False)
partition = recompute_ctx.get('partition', False)
logger.info(
"Start Recompute for PipeLineParallel. recompute_offload: {}, recompute_partition: {}"
.format(offload, partition))
"Start Recompute for PipeLineParallel. recompute_offload: {}, recompute_partition: {}".format(
offload, partition
)
)
world_size = dist.get_world_size()
self.global_rank = dist.get_rank()
......@@ -328,22 +354,28 @@ class PipelineLayer(Layer):
self._stage_id = self._topo.get_coord(self.global_rank).pipe
self._num_stages = self._topo.get_dim_size("pipe")
if num_stages:
assert self._num_stages == num_stages, "num_stages should be equal to be %d" % (
self._num_stages)
assert (
self._num_stages == num_stages
), "num_stages should be equal to be %d" % (self._num_stages)
else:
# construct default topology
if world_size % num_stages != 0:
raise ValueError(
"should provide correct num_stages({}) "
"which can be divided by world_size({})".format(
num_stages, world_size))
num_stages, world_size
)
)
dp_num = world_size // num_stages
self._topo = fleet.CommunicateTopology(["data", "pipe", "model"],
[dp_num, num_stages, 1])
self._topo = fleet.CommunicateTopology(
["data", "pipe", "model"], [dp_num, num_stages, 1]
)
self._stage_id = self._topo.get_coord(self.global_rank).pipe
self._num_stages = self._topo.get_dim_size("pipe")
self._total_stages_with_virtual_stages = self._num_stages * self._num_virtual_pipeline_stages
self._total_stages_with_virtual_stages = (
self._num_stages * self._num_virtual_pipeline_stages
)
# initialize segment
self._layers_desc = list(self.layers)
......@@ -381,16 +413,22 @@ class PipelineLayer(Layer):
start_idx = virtual_pp_rank * self._num_stages
for stage in range(self._num_stages):
# stage mark the real pp stage
if self.segment_parts[start_idx +
stage] <= layer_idx < self.segment_parts[
start_idx + stage + 1]:
if (
self.segment_parts[start_idx + stage]
<= layer_idx
< self.segment_parts[start_idx + stage + 1]
):
return stage
def get_num_virtual_stages(self):
return self._num_virtual_pipeline_stages
def get_model_chunks(self):
return None if self._num_virtual_pipeline_stages == 1 else self._model_chunks
return (
None
if self._num_virtual_pipeline_stages == 1
else self._model_chunks
)
def _construct_shared_comm(self):
shared_comm = {}
......@@ -398,17 +436,21 @@ class PipelineLayer(Layer):
return
layers_desc = self._layers_desc
shared_layer_names = set(s.layer_name for s in layers_desc
if isinstance(s, SharedLayerDesc))
shared_layer_names = set(
s.layer_name for s in layers_desc if isinstance(s, SharedLayerDesc)
)
for key in shared_layer_names:
shared_layers = []
for idx, layer in enumerate(layers_desc):
if isinstance(layer,
SharedLayerDesc) and layer.layer_name == key:
if (
isinstance(layer, SharedLayerDesc)
and layer.layer_name == key
):
shared_layers.append(idx)
shared_stages = set(
self.get_stage_from_index(idx) for idx in shared_layers)
self.get_stage_from_index(idx) for idx in shared_layers
)
self._dp_degree = self._topo.get_dim('data')
self._mp_degree = self._topo.get_dim('model')
self._sharding_degree = self._topo.get_dim('sharding')
......@@ -425,7 +467,9 @@ class PipelineLayer(Layer):
pipe=s,
data=dp,
sharding=sharding,
model=mp))
model=mp,
)
)
group = paddle.distributed.new_group(ranks=shared_ranks)
if self.global_rank in shared_ranks:
......@@ -434,8 +478,9 @@ class PipelineLayer(Layer):
shared_comm[key] = {
'ranks': shared_ranks,
'group': group,
'weight_attr':
self.shared_weight_attrs[key],
'weight_attr': self.shared_weight_attrs[
key
],
'layer': self.shared_layers[key],
}
return shared_comm
......@@ -443,10 +488,11 @@ class PipelineLayer(Layer):
def _synchronize_shared_weights(self):
for key, comm in self.shared_comm.items():
with paddle.framework.no_grad():
paddle.distributed.broadcast(getattr(comm['layer'],
comm['weight_attr']),
paddle.distributed.broadcast(
getattr(comm['layer'], comm['weight_attr']),
src=min(comm['ranks']),
group=comm['group'])
group=comm['group'],
)
for param in comm['layer'].parameters():
if self.global_rank != min(comm['ranks']):
......@@ -458,8 +504,9 @@ class PipelineLayer(Layer):
# need use trace_op to allreduce weight
if in_dygraph_mode():
with paddle.framework.no_grad():
paddle.distributed.all_reduce(param.grad,
group=comm['group'])
paddle.distributed.all_reduce(
param.grad, group=comm['group']
)
else:
with paddle.framework.no_grad():
paddle.fluid.framework._dygraph_tracer().trace_op(
......@@ -468,8 +515,9 @@ class PipelineLayer(Layer):
outputs={'Out': param._grad_ivar()},
attrs={
'ring_id': comm['group'].id,
'use_calc_stream': True
})
'use_calc_stream': True,
},
)
def _segment_network_for_interleave(self, seg_method):
logger.info("start segment network for interleave scheduler")
......@@ -477,14 +525,20 @@ class PipelineLayer(Layer):
self._layers_desc,
num_parts=self._num_stages,
method=seg_method,
num_virtual_pipeline_stage=self._num_virtual_pipeline_stages)
num_virtual_pipeline_stage=self._num_virtual_pipeline_stages,
)
self.segment_parts = seg.do_segment()
logger.info("segment result:" +
", ".join(str(arg) for arg in self.segment_parts))
logger.info(
"segment result:"
+ ", ".join(str(arg) for arg in self.segment_parts)
)
for i in range(self._stage_id, self._total_stages_with_virtual_stages,
self._num_stages):
for i in range(
self._stage_id,
self._total_stages_with_virtual_stages,
self._num_stages,
):
# If there are 2 real pp stages and 2 virtual pp stages, and the model has 8 layers.
# Layers [0, 1], [4, 5] will be assigned to the first real pp stage.
# Layers [2, 3], [6, 7] will be assigned to the second real pp stage.
......@@ -500,13 +554,15 @@ class PipelineLayer(Layer):
def _segment_network(self, seg_method):
logger.info("start segment network..")
seg = SegmentLayers(self._layers_desc,
num_parts=self._num_stages,
method=seg_method)
seg = SegmentLayers(
self._layers_desc, num_parts=self._num_stages, method=seg_method
)
self.segment_parts = seg.do_segment()
logger.info("segment result:" +
", ".join(str(arg) for arg in self.segment_parts))
logger.info(
"segment result:"
+ ", ".join(str(arg) for arg in self.segment_parts)
)
self._start_pos = self.segment_parts[self._stage_id]
self._end_pos = self.segment_parts[self._stage_id + 1]
......@@ -514,22 +570,30 @@ class PipelineLayer(Layer):
def _print_segmentation_for_debug(self):
# print information for debug
for stage in range(self._num_stages *
self._num_virtual_pipeline_stages):
for stage in range(
self._num_stages * self._num_virtual_pipeline_stages
):
start = self.segment_parts[stage]
end = self.segment_parts[stage + 1]
logger.info("stage={}, global_rank={} ,layer_number={}".format(
stage, self.global_rank, end - start))
logger.info(
"stage={}, global_rank={} ,layer_number={}".format(
stage, self.global_rank, end - start
)
)
for index, layer in enumerate(self._layers_desc[start:end]):
logger.info("{}: {}".format(index + start, str(layer)))
if self._num_virtual_pipeline_stages > 1:
for stage in range(self._num_stages):
stage_to_virtual_stage_info = "stage {} contains virtual stages: ".format(
stage)
for i in range(stage, self._total_stages_with_virtual_stages,
self._num_stages):
stage_to_virtual_stage_info = (
"stage {} contains virtual stages: ".format(stage)
)
for i in range(
stage,
self._total_stages_with_virtual_stages,
self._num_stages,
):
stage_to_virtual_stage_info += " {},".format(i)
logger.info(stage_to_virtual_stage_info)
......@@ -575,9 +639,11 @@ class PipelineLayer(Layer):
if layer.layer_name not in self.shared_layers:
self.shared_layers[layer.layer_name] = layer.build_layer()
self.shared_weight_attrs[
layer.layer_name] = layer.shared_weight_attr
layer.layer_name
] = layer.shared_weight_attr
for param in self.shared_layers[
layer.layer_name].parameters():
layer.layer_name
].parameters():
setattr(param, "is_firstly_shared", True)
if layer.forward_func is None:
......@@ -585,8 +651,11 @@ class PipelineLayer(Layer):
else:
run_function.append(
partial(layer.forward_func,
self.shared_layers[layer.layer_name]))
partial(
layer.forward_func,
self.shared_layers[layer.layer_name],
)
)
elif isinstance(layer, LayerDesc):
model = layer.build_layer()
......@@ -615,11 +684,15 @@ class PipelineLayer(Layer):
def forward(self, input, chunk_id=None):
if chunk_id is not None:
assert isinstance(chunk_id, int), "chunk_id should be an int"
assert self._num_virtual_pipeline_stages > 1, \
"chunk_id is only valid when using virtual pipeline stage"
assert chunk_id < len(self._model_chunks), \
"The virtual pipeline only has {} chunks, " \
"but received chunk_id {}.".format(len(self._model_chunks), chunk_id)
assert (
self._num_virtual_pipeline_stages > 1
), "chunk_id is only valid when using virtual pipeline stage"
assert chunk_id < len(self._model_chunks), (
"The virtual pipeline only has {} chunks, "
"but received chunk_id {}.".format(
len(self._model_chunks), chunk_id
)
)
# Get the target model chunk.
model_chunk = self._model_chunks[chunk_id]
# Update the self.run_function to the target run functions.
......@@ -637,20 +710,25 @@ class PipelineLayer(Layer):
funcs = self.run_function[start_idx:end_idx]
if not isinstance(input, tuple):
input = (input, )
input = (input,)
if self._need_recompute(funcs, input):
input = recompute_hybrid(
self.recompute_ctx,
self.forward_function(start_idx, end_idx), *input)
self.forward_function(start_idx, end_idx),
*input
)
else:
input = self.forward_function(start_idx, end_idx)(*input)
return input
def _need_recompute(self, funcs, inputs):
if not any(input_.stop_gradient == False
for input_ in inputs if isinstance(input_, paddle.Tensor)):
if not any(
input_.stop_gradient == False
for input_ in inputs
if isinstance(input_, paddle.Tensor)
):
return False
params = [f.parameters() for f in funcs if isinstance(f, Layer)]
......@@ -674,11 +752,18 @@ class PipelineLayer(Layer):
if self._num_virtual_pipeline_stages > 1:
# add virtual pipeline info to the save path
assert local_chunk_id is not None
virtual_pipeline_stage_message = "-virtual_pp_stage_{:0>2d}".format(
local_chunk_id)
layer_save_path = os.path.join(ckpt_dir,
'layer_{:0>2d}'.format(idx))
layer_save_path = layer_save_path + virtual_pipeline_stage_message + rank_message + '-model_states.pdparams'
virtual_pipeline_stage_message = (
"-virtual_pp_stage_{:0>2d}".format(local_chunk_id)
)
layer_save_path = os.path.join(
ckpt_dir, 'layer_{:0>2d}'.format(idx)
)
layer_save_path = (
layer_save_path
+ virtual_pipeline_stage_message
+ rank_message
+ '-model_states.pdparams'
)
return layer_save_path
def _save_model(run_functions, local_chunk_id=None):
......@@ -701,7 +786,8 @@ class PipelineLayer(Layer):
def set_state_dir(self, path):
assert os.path.exists(
path), "{} not found, please check the path".format(path)
path
), "{} not found, please check the path".format(path)
def _load_model(run_functions, local_chunk_id=None):
for idx, layer in enumerate(run_functions):
......@@ -715,21 +801,26 @@ class PipelineLayer(Layer):
pos_offset = self._start_poss[local_chunk_id]
layer_idx = idx + pos_offset
layer_save_path = os.path.join(
path, 'layer_{0:0>2d}'.format(layer_idx))
path, 'layer_{0:0>2d}'.format(layer_idx)
)
if self._num_virtual_pipeline_stages > 1:
# add virtual pipeline info to the path
assert local_chunk_id is not None
layer_save_path = layer_save_path + "-virtual_pp_stage_{:0>2d}".format(
local_chunk_id)
model_files = glob.glob(layer_save_path +
"*model_states.pdparams")
layer_save_path = (
layer_save_path
+ "-virtual_pp_stage_{:0>2d}".format(local_chunk_id)
)
model_files = glob.glob(
layer_save_path + "*model_states.pdparams"
)
model_files.sort()
mp_rank = self._topo.get_coord(self.global_rank).model
mp_world_size = self._topo.get_dim('model')
num_files = len(model_files)
load_param_path = model_files[mp_rank * num_files //
mp_world_size]
load_param_path = model_files[
mp_rank * num_files // mp_world_size
]
model_state_dict = paddle.load(load_param_path)
layer.set_state_dict(model_state_dict)
......
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或