wehub-resource-sync
420 lines
16 KiB
Python
420 lines
16 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import itertools
|
|
import logging
|
|
import re
|
|
from collections import OrderedDict
|
|
from enum import Enum
|
|
|
|
import paddle.distributed as dist
|
|
from paddle.distributed import fleet
|
|
from paddle.distributed.utils.log_utils import get_logger
|
|
|
|
from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
|
|
|
|
logger = get_logger("INFO", __name__)
|
|
|
|
|
|
class SplitPoint(Enum):
|
|
"""
|
|
Marking the position of the split.
|
|
BEGINNING: will split the model before the specified layer.
|
|
END: will split the model after the specified layer.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> import paddle.distributed as dist
|
|
|
|
>>> class MLP(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.fc1 = paddle.nn.Linear(8, 8)
|
|
... self.fc2 = paddle.nn.Linear(8, 8)
|
|
...
|
|
... def forward(self, input):
|
|
... return self.fc2(self.fc1(input))
|
|
|
|
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
|
>>> layer = MLP()
|
|
>>> pp_config = {
|
|
... 'fc1': dist.SplitPoint.END,
|
|
... }
|
|
"""
|
|
|
|
BEGINNING = 0
|
|
END = 1
|
|
|
|
|
|
class PipelineParallel(ParallelModel):
|
|
def __init__(self, model, split_spec, global_spec, pipeline_layers=None):
|
|
super().__init__(model)
|
|
self.split_spec = split_spec
|
|
self.global_spec = global_spec
|
|
self.pipeline_layers = pipeline_layers
|
|
self.pp_parallelizer = self.pipeline_parallel_fn
|
|
self.name_to_layer = {}
|
|
for layer_name, layer in model.named_sublayers():
|
|
self.name_to_layer[layer_name] = layer
|
|
|
|
def get_layer_by_name(self, name):
|
|
assert name in self.name_to_layer, (
|
|
f"layer name:{name} not in the model, please check the split_spec"
|
|
)
|
|
return self.name_to_layer[name]
|
|
|
|
def pipeline_parallel_fn(self, model):
|
|
mesh = fleet.auto.get_mesh()
|
|
pipeline_stage_num = mesh.get_dim_size("pp")
|
|
assert len(self.split_spec) == pipeline_stage_num - 1
|
|
|
|
def forward_post_hook(layer, input, output):
|
|
pipeline_stage_index = layer.pipeline_stage_index
|
|
split_point = layer.split_point
|
|
assert split_point == SplitPoint.END
|
|
# reshard to next pipeline stage
|
|
if isinstance(output, (dict, OrderedDict)):
|
|
for key, tensor in output.items():
|
|
assert is_tensor(tensor)
|
|
output[key] = dist.reshard(
|
|
tensor,
|
|
self.get_mesh(pipeline_stage_index + 1),
|
|
tensor.placements,
|
|
)
|
|
elif isinstance(output, list):
|
|
for i in range(len(output)):
|
|
assert is_tensor(output[i])
|
|
output[i] = dist.reshard(
|
|
output[i],
|
|
self.get_mesh(pipeline_stage_index + 1),
|
|
output[i].placements,
|
|
)
|
|
elif isinstance(output, tuple):
|
|
output = list(output)
|
|
for i in range(len(output)):
|
|
assert is_tensor(output[i])
|
|
output[i] = dist.reshard(
|
|
output[i],
|
|
self.get_mesh(pipeline_stage_index + 1),
|
|
output[i].placements,
|
|
)
|
|
output = tuple(output)
|
|
elif is_tensor(output):
|
|
output = dist.reshard(
|
|
output,
|
|
self.get_mesh(pipeline_stage_index + 1),
|
|
output.placements,
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"output between pp stages should be a dict of tensors or list of tensors or tuple of tensors or tensor, but {type(output)}"
|
|
)
|
|
return output
|
|
|
|
def forward_pre_hook(layer, input):
|
|
split_point = layer.split_point
|
|
assert split_point == SplitPoint.BEGINNING
|
|
# TODO(deepllz): support in the future
|
|
return input
|
|
|
|
# step1: set every layer's own pipeline_stage_index
|
|
split_layer_names = list(self.split_spec.keys())
|
|
sublayer_names = [name for name, _ in model.named_sublayers()]
|
|
# Mark which layer is the next pipeline stage
|
|
pipeline_layer_mark = [0 for _ in range(len(sublayer_names))]
|
|
for split_layer_name in split_layer_names:
|
|
split_point = self.split_spec[split_layer_name]
|
|
index = sublayer_names.index(split_layer_name)
|
|
if split_point == SplitPoint.END:
|
|
is_valid = False
|
|
for i in range(index + 1, len(sublayer_names)):
|
|
if not sublayer_names[i].startswith(split_layer_name):
|
|
pipeline_layer_mark[i] = 1
|
|
is_valid = True
|
|
break
|
|
assert is_valid, (
|
|
f"the last layer:{split_layer_name} must not be SplitPoint.END, please check the split_spec"
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
"SplitPoint.BEGINNING is not supported currently"
|
|
)
|
|
pipeline_layer_mark[index] = 1
|
|
# the inclusiveSum of pipeline_layer_mark is the pipeline stage index
|
|
pipeline_stage_index = list(itertools.accumulate(pipeline_layer_mark))
|
|
for index, (name, layer) in enumerate(model.named_sublayers()):
|
|
layer.pipeline_stage_index = pipeline_stage_index[index]
|
|
|
|
# step2: insert reshard
|
|
for name in split_layer_names:
|
|
layer = self.get_layer_by_name(name)
|
|
split_point = self.split_spec[name]
|
|
layer.split_point = split_point
|
|
if split_point == SplitPoint.END:
|
|
layer.register_forward_post_hook(forward_post_hook)
|
|
else:
|
|
raise NotImplementedError(
|
|
"SplitPoint.BEGINNING is not supported currently"
|
|
)
|
|
layer.register_forward_pre_hook(forward_pre_hook)
|
|
|
|
if self.global_spec:
|
|
self.process_global_mesh_layers()
|
|
|
|
return model
|
|
|
|
def process_global_mesh_layers(self):
|
|
g_mesh = fleet.auto.get_mesh()
|
|
g_mesh = g_mesh.get_mesh_with_dim("pp")
|
|
|
|
def forward_post_hook(layer, input, output):
|
|
if isinstance(output, (list, tuple)):
|
|
global_output = list(output)
|
|
for ind in range(len(global_output)):
|
|
output_i = global_output[ind]
|
|
if is_tensor(output_i):
|
|
if output_i.is_dist():
|
|
global_output[ind] = dist.reshard(
|
|
output_i,
|
|
g_mesh,
|
|
[
|
|
dist.Replicate()
|
|
for _ in range(len(g_mesh._shape))
|
|
],
|
|
)
|
|
else:
|
|
global_output[ind] = dist.shard_tensor(
|
|
output_i,
|
|
g_mesh,
|
|
[
|
|
dist.Replicate()
|
|
for _ in range(len(g_mesh._shape))
|
|
],
|
|
)
|
|
|
|
if isinstance(output, tuple):
|
|
global_output = tuple(global_output)
|
|
return global_output
|
|
elif is_tensor(output):
|
|
if output.is_dist():
|
|
return dist.reshard(
|
|
output,
|
|
g_mesh,
|
|
[dist.Replicate() for _ in range(len(g_mesh._shape))],
|
|
)
|
|
else:
|
|
return dist.shard_tensor(
|
|
output,
|
|
g_mesh,
|
|
[dist.Replicate() for _ in range(len(g_mesh._shape))],
|
|
)
|
|
else:
|
|
raise TypeError(
|
|
"layer output can only be tensor or list/tuple of tensor"
|
|
)
|
|
|
|
def forward_pre_hook(layer, args, kwargs):
|
|
pp_idx = getattr(layer, "pipeline_stage_index", 0)
|
|
new_args = []
|
|
new_kwargs = {}
|
|
|
|
def reshard_not_mesh_match_tensor(arg):
|
|
cur_pp_mesh = self.get_mesh(pp_idx)
|
|
if (
|
|
arg is not None
|
|
and is_tensor(arg)
|
|
and arg.is_dist()
|
|
and arg.process_mesh != cur_pp_mesh
|
|
):
|
|
return dist.reshard(
|
|
arg,
|
|
cur_pp_mesh,
|
|
[dist.Replicate(), dist.Replicate()],
|
|
)
|
|
return arg
|
|
|
|
for arg in args:
|
|
new_args.append(reshard_not_mesh_match_tensor(arg))
|
|
|
|
for key, arg in kwargs.items():
|
|
new_kwargs[key] = reshard_not_mesh_match_tensor(arg)
|
|
|
|
return (tuple(new_args), new_kwargs)
|
|
|
|
# wa because of pir in vpp mode send receive bug
|
|
for layer_name in self.global_spec:
|
|
layer = self.get_layer_by_name(layer_name)
|
|
layer.register_forward_post_hook(forward_post_hook)
|
|
|
|
if self.pipeline_layers is not None:
|
|
for layer_name in self.pipeline_layers:
|
|
layer = self.get_layer_by_name(layer_name)
|
|
layer.register_forward_pre_hook(
|
|
forward_pre_hook, with_kwargs=True
|
|
)
|
|
else:
|
|
for layer in self.name_to_layer.values():
|
|
layer.register_forward_pre_hook(
|
|
forward_pre_hook, with_kwargs=True
|
|
)
|
|
|
|
|
|
def pipeline_parallel(model, optimizer=None, config=None):
|
|
"""
|
|
pipeline_parallel converts model and optimizer to pipelined distributed model
|
|
|
|
Args:
|
|
model (paddle.nn.Layer): A single card model to be distributed
|
|
optimizer (paddle.optimizer.Optimizer): An optimizer to be distributed
|
|
config (dict): {
|
|
"split_spec": OrderedDict|dict|str|list(str), The pipeline parallel split point.
|
|
if split_spec is a string or list, such as "llama.layer" or ["llama.layerA", "llama.layerB"], Then the layer with same prefix a will be divided equally according to the size of pipeline degree.
|
|
if split_spec is a OrderedDict|dict, key is the layer name, and the value is the split position that can be SplitPoint.BEGINNING or SplitPoint.END, the order of the keys is the order of the pipeline stage.
|
|
NOTE: dict is also ordered after python3.7, so use dict at this time.
|
|
"global_spec": str|list(str), make the output tensor of specific layers on global mesh.
|
|
}
|
|
|
|
Returns:
|
|
PipelineParallel: a distributed model
|
|
ParallelOptimizer: a distributed optimizer
|
|
"""
|
|
|
|
split_spec = config.get("split_spec")
|
|
if split_spec is None:
|
|
logging.warning("No split_spec, pipeline parallel won't do anything.")
|
|
return model, optimizer
|
|
|
|
mesh = fleet.auto.get_mesh()
|
|
assert mesh is not None, (
|
|
"global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
|
|
)
|
|
assert "pp" in mesh.dim_names, (
|
|
"pp must in the mesh dim_names when use pipeline_parallel"
|
|
)
|
|
|
|
global_spec = config.get("global_spec")
|
|
if isinstance(split_spec, str):
|
|
split_spec = [split_spec]
|
|
|
|
matched_layer_name = None
|
|
if isinstance(split_spec, (list, tuple)):
|
|
# match layer_name with split_spec following by a dot and numbers and no other characters
|
|
# such as split_spec = ["llama.layer"], then llama.layer.0 is matched, llama.layer.0.mlp is not matched
|
|
patterns = [rf"{prefix}\.\d+$" for prefix in split_spec]
|
|
|
|
def is_match(layer_name):
|
|
for pattern in patterns:
|
|
if re.match(pattern, layer_name) or layer_name in split_spec:
|
|
return True
|
|
return False
|
|
|
|
def filter_matched_layer(matched_layer_name):
|
|
# remove the base name if it has a numbered suffix
|
|
string_set = set(matched_layer_name)
|
|
to_remove = set()
|
|
|
|
numbered_pattern = re.compile(r'^(.+)\.\d+$')
|
|
for s in matched_layer_name:
|
|
match = numbered_pattern.match(s)
|
|
if match:
|
|
base_name = match.group(1)
|
|
if base_name in string_set:
|
|
to_remove.add(base_name)
|
|
|
|
res = []
|
|
for s in matched_layer_name:
|
|
if s not in to_remove:
|
|
res.append(s)
|
|
return res
|
|
|
|
matched_layer_name = [
|
|
name for name, _ in model.named_sublayers() if is_match(name)
|
|
]
|
|
matched_layer_name = filter_matched_layer(matched_layer_name)
|
|
pp_size = mesh.get_dim_size("pp")
|
|
layer_num = len(matched_layer_name)
|
|
assert layer_num > 0, (
|
|
"No layer match the split_spec, please check its correctness"
|
|
)
|
|
assert layer_num >= pp_size, (
|
|
"The number of layers must not be less than the pp size"
|
|
)
|
|
if layer_num % pp_size != 0:
|
|
logger.warning(
|
|
f"The number of layers({layer_num}) must be divisible by the pp size({pp_size}), but got {layer_num} and {pp_size}"
|
|
)
|
|
|
|
def divide_list_indices(n, k):
|
|
base_size = n // k
|
|
extra = n % k
|
|
|
|
indices = []
|
|
current_index = -1
|
|
|
|
for i in range(k - 1):
|
|
current_index += base_size
|
|
if i < extra:
|
|
current_index += 1
|
|
indices.append(current_index)
|
|
return indices
|
|
|
|
indices = divide_list_indices(layer_num, pp_size)
|
|
split_spec_dict = OrderedDict(
|
|
[
|
|
(matched_layer_name[indices[i]], SplitPoint.END)
|
|
for i in range(pp_size - 1)
|
|
]
|
|
)
|
|
else:
|
|
layers_per_rank = layer_num // pp_size
|
|
split_spec_dict = OrderedDict(
|
|
[
|
|
(
|
|
matched_layer_name[i * layers_per_rank - 1],
|
|
SplitPoint.END,
|
|
)
|
|
for i in range(1, pp_size)
|
|
]
|
|
)
|
|
else:
|
|
sublayer_names = [name for name, _ in model.named_sublayers()]
|
|
split_spec_dict = split_spec
|
|
for key, value in split_spec_dict.items():
|
|
assert key in sublayer_names, (
|
|
f"wrong split layer, expected one of {sublayer_names}"
|
|
)
|
|
assert value is SplitPoint.END, "not supported split point at now."
|
|
|
|
if global_spec:
|
|
if isinstance(global_spec, str):
|
|
global_spec = [global_spec]
|
|
else:
|
|
assert isinstance(global_spec, (list, tuple)), (
|
|
f"global_spec can only be list or list(str), but got:{type(global_spec)}"
|
|
)
|
|
|
|
logger.info(
|
|
f"split_spec_dict: {split_spec_dict}, global_spec: {global_spec}, matched_layer_name: {matched_layer_name}"
|
|
)
|
|
|
|
model = PipelineParallel(
|
|
model, split_spec_dict, global_spec, matched_layer_name
|
|
)
|
|
if optimizer is not None:
|
|
optimizer = ParallelOptimizer(optimizer)
|
|
|
|
return model, optimizer
|