wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/paddle/distributed/auto_parallel/intermediate/__init__.py
+15
View File
@@ -0,0 +1,15 @@
# 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.
__all__ = []
python/paddle/distributed/auto_parallel/intermediate/context_parallel.py
+391
View File
@@ -0,0 +1,391 @@
# 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.
from __future__ import annotations
import logging
import paddle
import paddle.distributed as dist
from paddle.distributed.auto_parallel.ring_attention import (
shard_seq_load_balance,
)
from .tensor_parallel import PlanBase
class PrepareContextParallel(PlanBase):
"""
Prepare Input for context parallel optimizations.
This will work for Layer that calls like whole-llama Layer which is the first layer in the network.
Users can set backend='p2p/all2all' for different context parallel strategys.
backend='p2p' will use Ring FlashAttention strategy which segments input with balance in the sequence dimension before whole-llama Layer.
backend='all2all' will use Deepspeed Ulysses strategy(Paddle SegmentParallel strategy) which segments input in the sequence dimension before whole-llama Layer.
Args:
backend (string): select strategy for context parallel, now support 'p2p' and 'all2all'.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.distributed as dist
>>> class SDPALayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
...
... def forward(self, q, k, v):
... return paddle.nn.functional.scaled_dot_product_attention(q, k, v)
>>>
>>> class AttentionLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.hidden_size = 64
... self.num_key_value_heads = 10
... self.head_dim = 64
... self.sdpa = SDPALayer()
... self.q = paddle.nn.Linear(
... self.hidden_size,
... self.hidden_size,
... bias_attr=False,
... )
... self.k = paddle.nn.Linear(
... self.hidden_size,
... self.num_key_value_heads * self.head_dim,
... bias_attr=False,
... )
... self.v = paddle.nn.Linear(
... self.hidden_size,
... self.num_key_value_heads * self.head_dim,
... bias_attr=False,
... )
...
... def forward(self, input):
... q = self.q(input)
... k = self.k(input)
... v = self.v(input)
... return self.sdpa(q, k, v)
>>>
>>> class LlamaLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.attention = AttentionLayer()
...
... def forward(self, input, label):
... return self.attention(input)
>>>
>>> class LlamaForCausalLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.llama = LlamaLayer()
... self.weight = self.create_parameter(shape=[64, 1024])
... self.loss_func = paddle.nn.CrossEntropyLoss()
...
... def forward(self, input, label):
... out = self.llama(input, label)
... logits = paddle.matmul(out, self.weight)
... loss = self.loss_func(logits, label)
... return logits
>>>
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
>>> layer = LlamaForCausalLayer()
>>> mp_config = {
... 'llama': dist.PrepareContextParallel('p2p'),
... 'sdpa': dist.ContextParallel('p2p'),
... }
"""
def __init__(self, backend: str = 'p2p') -> None:
super().__init__()
self.backend = backend
assert self.backend in [
'p2p',
'all2all',
], f"backend must be 'p2p' or 'all2all', but got {self.backend}"
def all2all_split_input_pre_hook(self, process_mesh):
def shard_tensor(input_tensor, seq_dim):
cp_index = process_mesh.dim_names.index('sep')
placements = input_tensor.placements
if placements is None:
placements = [
dist.Replicate() for _ in range(len(process_mesh.shape))
]
# split sequence dim
placements[cp_index] = dist.Shard(seq_dim)
reshard_input = dist.reshard(input_tensor, process_mesh, placements)
return reshard_input
def all2all_split_input(layer, args):
cp_index = process_mesh.dim_names.index('sep')
cp_degree = process_mesh.shape[cp_index]
# check input_ids
if isinstance(args, (list, tuple)):
all_args = []
for input_tensor in args:
assert input_tensor.is_dist(), (
"Input tensor must be a distributed tensor."
)
assert len(input_tensor.shape) == 2, (
f"input_ids should be [batch_size, seq_len], but got {input_tensor.shape}"
)
_, seq_len = input_tensor.shape
assert seq_len % cp_degree == 0, (
f"sequence length {seq_len} must be divisible by cp degree {cp_degree}"
)
reshard_input = shard_tensor(input_tensor, 1)
all_args.append(reshard_input)
new_args = tuple(all_args)
return new_args
elif isinstance(args, paddle.Tensor):
reshard_input = shard_tensor(args, 1)
return reshard_input
else:
raise ValueError(
f"Unsupported argument type: {type(args)}. Expected list of tensors or single tensor."
)
return all2all_split_input
def p2p_split_input_pre_hook(self, process_mesh):
def p2p_split_input(layer, args):
cp_index = process_mesh.dim_names.index('sep')
cp_degree = process_mesh.shape[cp_index]
if isinstance(args, (list, tuple)):
all_args = []
for input_tensor in args:
# check input_ids
assert input_tensor.is_dist(), (
"Input tensor must be a distributed tensor."
)
assert len(input_tensor.shape) == 2, (
f"input_ids should be [batch_size, seq_len], but got {input_tensor.shape}"
)
placements = input_tensor.placements
if placements is None:
placements = [
dist.Replicate()
for _ in range(len(process_mesh.shape))
]
assert placements[cp_index] == dist.Replicate(), (
"Input tensor must be a replicated tensor in cp mesh."
)
reshard_input = shard_seq_load_balance(input_tensor, 1)
all_args.append(reshard_input)
new_args = tuple(all_args)
return new_args
elif isinstance(args, paddle.Tensor):
reshard_input = shard_seq_load_balance(input_tensor, 1)
return reshard_input
else:
raise ValueError(
f"Unsupported argument type: {type(args)}. Expected list of tensors or single tensor."
)
return p2p_split_input
def apply(self, layer, process_mesh, shard_param_list):
if self.backend == 'all2all':
# Deepspeed Ulysses
layer.register_forward_pre_hook(
self.all2all_split_input_pre_hook(process_mesh)
)
elif self.backend == 'p2p':
# Ring FlashAttention
layer.register_forward_pre_hook(
self.p2p_split_input_pre_hook(process_mesh)
)
else:
logging.warning(
f'{self.backend} is not supported backend for context parallel'
)
class ContextParallel(PlanBase):
"""
Applies context parallel optimizations to the attention layer.
This will work for Layer that calls paddle.nn.functional.scaled_dot_product_attention).
Users can set backend='p2p/all2all' for different context parallel strategys.
backend='p2p' will use Ring FlashAttention strategy which segments q/k/v in the sequence dimension and communicates k/v between ranks.
backend='all2all' will use Deepspeed Ulysses strategy(Paddle SegmentParallel strategy) which inserts all2all before and after sdpa compute.
Note:
Args:
backend (string): select strategy for context parallel, now support 'p2p' and 'all2all'.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.distributed as dist
>>> class SDPALayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
...
... def forward(self, q, k, v):
... return paddle.nn.functional.scaled_dot_product_attention(q, k, v)
>>>
>>> class AttentionLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.hidden_size = 64
... self.num_key_value_heads = 10
... self.head_dim = 64
... self.sdpa = SDPALayer()
... self.q = paddle.nn.Linear(
... self.hidden_size,
... self.hidden_size,
... bias_attr=False,
... )
... self.k = paddle.nn.Linear(
... self.hidden_size,
... self.num_key_value_heads * self.head_dim,
... bias_attr=False,
... )
... self.v = paddle.nn.Linear(
... self.hidden_size,
... self.num_key_value_heads * self.head_dim,
... bias_attr=False,
... )
...
... def forward(self, input):
... q = self.q(input)
... k = self.k(input)
... v = self.v(input)
... return self.sdpa(q, k, v)
>>>
>>> class LlamaLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.attention = AttentionLayer()
...
... def forward(self, input, label):
... return self.attention(input)
>>>
>>> class LlamaForCausalLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.llama = LlamaLayer()
... self.weight = self.create_parameter(shape=[64, 1024])
... self.loss_func = paddle.nn.CrossEntropyLoss()
...
... def forward(self, input, label):
... out = self.llama(input, label)
... logits = paddle.matmul(out, self.weight)
... loss = self.loss_func(logits, label)
... return logits
>>>
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
>>> layer = LlamaForCausalLayer()
>>> mp_config = {
... 'llama': dist.PrepareContextParallel('p2p'),
... 'sdpa': dist.ContextParallel('p2p'),
... }
"""
def __init__(self, backend: str = 'p2p') -> None:
super().__init__()
self.backend = backend
def all2all_reshard_pre_hook(self, process_mesh):
def all2all_reshard_hook(layer, args):
cp_index = process_mesh.dim_names.index('sep')
cp_degree = process_mesh.shape[cp_index]
all_args = []
for arg in args:
# check q k v
assert arg.is_dist(), f"arg {arg} must be a distributed tensor."
assert len(arg.shape) == 3 or len(arg.shape) == 4
placements = arg.placements
assert placements[cp_index] == dist.Shard(1), (
f"arg {arg} must be sharded in sequence dimension."
)
# reshard [batch_size，seq_len/sep，num_head，head_dim] -> [batch_size，seq_len，num_head/sep，head_dim]
placements[cp_index] = dist.Shard(2)
target_arg = dist.reshard(arg, process_mesh, placements)
all_args.append(target_arg)
new_args = tuple(all_args)
return new_args
return all2all_reshard_hook
def all2all_reshard_post_hook(self, process_mesh):
def all2all_reshard_hook(layer, input, output):
cp_index = process_mesh.dim_names.index('sep')
cp_degree = process_mesh.shape[cp_index]
placements = output.placements
assert output.is_dist(), (
f"output {output} must be a distributed tensor."
)
assert len(output.shape) == 4 or len(output.shape) == 3
assert placements[cp_index] == dist.Shard(2), (
f"output {output} must be Shard(2) in sequence dimension."
)
# reshard [batch_size，seq_len，num_head/seq，head_dim] -> [batch_size，seq_len/sep，num_head，head_dim]
placements[cp_index] = dist.Shard(1)
target_output = dist.reshard(output, process_mesh, placements)
return target_output
return all2all_reshard_hook
def p2p_reshard_pre_hook(self, process_mesh):
def input_hook(layer, args, kwargs):
cp_index = process_mesh.dim_names.index('sep')
cp_degree = process_mesh.shape[cp_index]
for arg in args:
# check q k v
assert arg.is_dist(), (
"Input tensor must be a distributed tensor."
)
assert len(arg.shape) == 3 or len(arg.shape) == 4
placements = arg.placements
assert placements[cp_index] == dist.Shard(1), (
f"arg {arg} must be Shard(1) in sequence dimension."
)
# edit kwarg backend to 'p2p'
new_kwargs = kwargs
new_kwargs['backend'] = 'p2p'
return args, new_kwargs
return input_hook
def apply(self, layer, process_mesh, shard_param_list):
if self.backend == 'all2all':
# Deepspeed Ulysses
layer.register_forward_pre_hook(
self.all2all_reshard_pre_hook(process_mesh)
)
layer.register_forward_post_hook(
self.all2all_reshard_post_hook(process_mesh)
)
elif self.backend == 'p2p':
# Ring FlashAttention
layer.register_forward_pre_hook(
self.p2p_reshard_pre_hook(process_mesh), with_kwargs=True
)
else:
logging.warning(
f'{self.backend} is not supported backend for context parallel'
)
python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
+294
View File
@@ -0,0 +1,294 @@
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import paddle
import paddle.distributed as dist
from paddle import pir
from paddle.base.framework import (
in_dygraph_mode,
in_pir_mode,
)
from paddle.distributed import fleet
from paddle.nn import Layer
from paddle.optimizer import Optimizer
def is_tensor(tensor):
if in_dygraph_mode():
return isinstance(tensor, paddle.Tensor)
elif in_pir_mode():
return isinstance(tensor, pir.Value)
else:
raise RuntimeError(
"PipelineParallel are only supported in dynamic or pir mode."
)
class ParallelOptimizer:
def __init__(
self,
optimizer,
level=None,
sharding_mesh_dim=None,
):
self.level = None
self.sharding_mesh_dim = None
self.optimizer = None
if isinstance(optimizer, ParallelOptimizer):
self.optimizer = optimizer.optimizer
if level is None:
self.level = optimizer.level
self.sharding_mesh_dim = optimizer.sharding_mesh_dim
else:
if isinstance(level, int):
level = str(level)
assert level in ("0", "1", "2", "3", None)
if optimizer.level is not None:
assert level == optimizer.level, (
f"The level passed in is not identical with previous level. Current level is {level}, previous level is {optimizer.level}"
)
self.level = level
self.sharding_mesh_dim = sharding_mesh_dim
else:
assert isinstance(optimizer, Optimizer)
self.optimizer = optimizer
if isinstance(level, int):
level = str(level)
assert level in ("0", "1", "2", "3", None)
# level=0 and level=None are all mean pure dp
self.level = level
self.sharding_mesh_dim = sharding_mesh_dim
self.is_initialized = False
def parallelize(self):
assert self.optimizer is not None
if self.is_initialized:
return self.optimizer
mesh = fleet.auto.get_mesh()
if self.level == "1":
self.optimizer = dist.shard_optimizer(
self.optimizer,
dist.ShardingStage1(self.sharding_mesh_dim, mesh),
)
elif self.level == "2":
self.optimizer = dist.shard_optimizer(
self.optimizer,
dist.ShardingStage2(self.sharding_mesh_dim, mesh),
)
elif self.level == "3":
self.optimizer = dist.shard_optimizer(
self.optimizer,
dist.ShardingStage3(self.sharding_mesh_dim, mesh),
)
else:
self.optimizer = dist.shard_optimizer(self.optimizer, None)
self.is_initialized = True
return self.optimizer
def update_param_list(self, parallelized_parameters):
self.optimizer._parameter_list = parallelized_parameters
if isinstance(parallelized_parameters[0], dict):
self.optimizer._param_groups = []
for param_group in self.parallelized_parameters:
self.optimizer._add_param_group(param_group.copy())
else:
self.optimizer._param_groups = self.optimizer._parameter_list
class ParallelModel:
def __init__(self, model):
super().__init__()
self.pp_parallelizer = None
self.tp_parallelizer = None
self.sharding_parallelizer = None
self.model = None
self.share_param_list = {}
self.layer_param_placements = {}
if isinstance(model, ParallelModel):
self.pp_parallelizer = model.pp_parallelizer
self.tp_parallelizer = model.tp_parallelizer
self.sharding_parallelizer = model.sharding_parallelizer
self.model = model.model
else:
assert isinstance(model, Layer)
self.model = model
self.is_parallelized = False
def get_mesh(self, pp_idx=0):
mesh = fleet.auto.get_mesh()
if "pp" in mesh.dim_names:
mesh = mesh.get_mesh_with_dim("pp", pp_idx)
return mesh
def parallelize_model(self):
assert self.model is not None
if self.is_parallelized:
return self.model
if self.pp_parallelizer is not None:
assert callable(self.pp_parallelizer)
self.model = self.pp_parallelizer(self.model)
if self.tp_parallelizer is not None:
assert callable(self.tp_parallelizer)
self.model, self.layer_param_placements = self.tp_parallelizer(
self.model
)
if self.sharding_parallelizer is not None:
assert callable(self.sharding_parallelizer)
self.model = self.sharding_parallelizer(self.model)
self._shard_all_param(self.model)
self.is_parallelized = True
return self.model
def _process_share_weight_layer(
self, layer, origin_weight, param_name, param_placements
):
ipp = (
layer.pipeline_stage_index
if hasattr(layer, "pipeline_stage_index")
else 0
)
def create_pre_hook(origin_weight, param_name):
def forward_pre_hook(layer, input):
setattr(
layer,
param_name,
None,
)
delattr(layer, param_name)
mesh = self.get_mesh(ipp)
share_weight = dist.reshard(
origin_weight,
mesh,
param_placements,
)
setattr(
layer,
param_name,
share_weight,
)
return forward_pre_hook
def create_post_hook(origin_weight, param_name):
def forward_post_hook(layer, input, output):
setattr(
layer,
param_name,
origin_weight,
)
return forward_post_hook
layer.register_forward_pre_hook(
create_pre_hook(origin_weight, param_name)
)
layer.register_forward_post_hook(
create_post_hook(origin_weight, param_name)
)
def _shard_all_param(self, model):
param_name_to_shard_param = {}
param_name_to_pp_stage = {}
def shard_layer_param(layer):
if self.pp_parallelizer is not None:
assert hasattr(layer, "pipeline_stage_index")
for param_name in list(layer._parameters.keys()):
param = getattr(layer, param_name)
if param is not None:
param_full_name = param.name
ipp = (
layer.pipeline_stage_index
if hasattr(layer, "pipeline_stage_index")
else 0
)
mesh = self.get_mesh(ipp)
param_placements = [
dist.Replicate() for _ in range(len(mesh._shape))
]
if layer in self.layer_param_placements:
if param_name in self.layer_param_placements[layer]:
param_placements = (
self.layer_param_placements[layer][param_name]
if self.layer_param_placements[layer][
param_name
]
is not None
else param_placements
)
if not param.is_dist():
if param_full_name in param_name_to_shard_param:
setattr(
layer,
param_name,
param_name_to_shard_param[param_full_name],
)
if ipp != param_name_to_pp_stage[param_full_name]:
self._process_share_weight_layer(
layer,
param_name_to_shard_param[param_full_name],
param_name,
param_placements,
)
else:
param = dist.shard_tensor(
param, mesh, param_placements
)
param_name_to_shard_param[param_full_name] = param
param_name_to_pp_stage[param_full_name] = ipp
setattr(layer, param_name, param)
else:
if (
param_full_name in param_name_to_shard_param
and ipp != param_name_to_pp_stage[param_full_name]
):
self._process_share_weight_layer(
layer,
param_name_to_shard_param[param_full_name],
param_name,
param_placements,
)
elif param_full_name not in param_name_to_shard_param:
param_name_to_shard_param[param_full_name] = param
param_name_to_pp_stage[param_full_name] = ipp
for name, layer in model.named_sublayers():
shard_layer_param(layer)
def parallelize_model_and_optimizer(model, optimizer=None):
if not isinstance(model, ParallelModel):
assert not isinstance(optimizer, ParallelOptimizer)
logging.warning(
"The method `parallelize_model_and_optimizer` won't do anything since the model is not parallelized."
)
return model, optimizer
parallelized_model = model.parallelize_model()
parallelized_optimizer = None
if optimizer is not None:
assert isinstance(optimizer, ParallelOptimizer)
optimizer.update_param_list(parallelized_model.parameters())
parallelized_optimizer = optimizer.parallelize()
return parallelized_model, parallelized_optimizer
python/paddle/distributed/auto_parallel/intermediate/parallelize.py
+385
View File
@@ -0,0 +1,385 @@
# 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.
from __future__ import annotations
import warnings
from typing import TYPE_CHECKING, TypedDict
from typing_extensions import NotRequired
from paddle.distributed import fleet
from paddle.framework import core
from .parallel_base import ParallelOptimizer, parallelize_model_and_optimizer
from .pipeline_parallel import pipeline_parallel
from .sharded_data_parallel import sharded_data_parallel
from .tensor_parallel import tensor_parallel
if TYPE_CHECKING:
import paddle
from .pipeline_parallel import SplitPoint
from .tensor_parallel import PlanBase
class _DPConfig(TypedDict):
sharding_level: str | int
class _MPConfig(TypedDict):
parallelize_plan: dict[str, PlanBase | list[PlanBase]]
class _PPConfig(TypedDict):
split_spec: str | dict[str, SplitPoint]
global_spec: NotRequired[str]
class _ParallelizeConfig(TypedDict):
dp_config: NotRequired[_DPConfig]
mp_config: NotRequired[_MPConfig]
pp_config: NotRequired[_PPConfig]
def parallelize(
model: paddle.nn.Layer,
optimizer: paddle.optimizer.Optimizer | None = None,
mesh: paddle.distributed.ProcessMesh | None = None,
config: _ParallelizeConfig | None = None,
) -> tuple[paddle.nn.Layer, paddle.optimizer.Optimizer]:
"""
Parallelize the model and optimizer from a single card version to a distributed version.
Args:
model (paddle.nn.Layer): the model to be parallelized.
optimizer (paddle.optimizer.Optimizer, optional): the optimizer to be parallelized.
Could be `None` if no optimizer to be parallelized.
mesh (paddle.distributed.ProcessMesh, optional): the process mesh for parallelize the model and the optimizer.
Best practice: calling `dist.auto_parallel.set_mesh` to set the global mesh ahead of calling `parallelize`
and keep the `mesh` parameter as `None.
If the `mesh` is not None, the mesh passed to `parallelize` will overwrite the mesh set by `set_mesh`.
config (dict, optional): a dict contains the parallel config.
The keys of the dict can be chosen from `dp_config`, `mp_config` and `pp_config` which will be used to
determine the parallel method for data parallel, tensor parallel and pipeline parallel separately.
A valid config can be like this: {"dp_config": for more information refer the `dp_config` section of
this doc, "mp_config": for more information refer the `mp_config` section of this doc, "pp_config":
for more information refer the `pp_config` section of this doc}.
dp_config (dict): a dict specifying the data parallel config. The keys of `dp_config` is `sharding_level`.
The value of `sharding_level` can be chosen from 0/1/2/3, which means pure data parallel, sharding
parallel stage 1, sharding parallel stage 2 and sharding parallel stage 3 separately. A valid
dp_config can be like this: {"sharding_level": 2}.
mp_config (dict): a dict specifying the tensor parallel config. The keys of `mp_config` is
`parallelize_plan`. The value of `parallelize_plan` is another dict, mapping a layer name or a param
name to a specific parallel plan. Note that the layer name could be written in regular format. If
mapping a param name to a specific plan, the name of the param must be ended with `weight` or `bias`.
And all valid parallel plan is `ColWiseParallel`, `RowWiseParallel`, `SequenceParallelBegin,
`SequenceParallelDisable`, `SequenceParallelEnable`, `SequenceParallelEnd`, `PrepareLayerInput` and
`PrepareLayerOutput`. A valid mp_config can be like this: {"llama.embed_tokens": dist.ColWiseParallel(),
"llama.norm": dist.SequenceParallelEnable(), "lm_head.weight": dist.ColWiseParallel()}.
pp_config (dict): a dict specifying the pipeline parallel config. The keys of `pp_config` is `split_spec`
and `global_spec`. The `split_spec` can be a dict or a string. If the `split_spec` is a dict, it maps
a layer name to a `SplitPoint`, note that the layer name could be written in regular format. The
pipeline parallel will exactly split the model at the point indicated by the map. If the `split_spec`
is a string, it contains the prefix of a set of layers. The pipeline parallel will automatically split
the model evenly at target layer. The `global_spec` is a string indicating a layer that contains global
tensors, which will be duplicated through all stages of the pipeline parallel. Some valid pp_config
can be list these: {"split_spec": "llama.layers", "global_spec": "llama.global_layer"}
or {"split_spec": {"llama.layers.1": SplitPoint.END}}.
cp_config (dict): a dict specifying the context parallel config. The keys of `cp_config` is
`parallelize_plan`. The value of `parallelize_plan` is another dict, mapping a layer name or a param
name to a specific parallel plan. All valid parallel plan is `ContextParallel` and `PrepareContextParallel`.
A valid cp_config can be like this: {"llama": dist.PrepareContextParallel('p2p'),
"llama.sdpa": dist.ContextParallel('p2p')}.
Note:
If the mesh is `None` or neither of `dp_config`, `mp_config`, `pp_config` and `cp_config` is in the config, this
api will do nothing but return the model and optimizer passed in.
Returns:
model, optimizer: the model and the optimizer after parallelize
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.distributed as dist
>>> class ModelConfig:
... def __init__(self):
... self.vocab_size = 10
... self.hidden_size = 20
... self.intermediate_size = 20
... self.num_layers = 2
>>> model_config = ModelConfig()
>>> class LlamaRMSNorm(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.weight = paddle.create_parameter(
... shape=[model_config.hidden_size],
... dtype=paddle.get_default_dtype(),
... )
...
... def forward(self, input):
... pass
>>> class LlamaAttention(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
...
... self.qkv_proj = paddle.nn.Linear(
... model_config.hidden_size,
... model_config.hidden_size * 3,
... bias_attr=False,
... )
...
... self.o_proj = paddle.nn.Linear(
... model_config.hidden_size,
... model_config.hidden_size,
... bias_attr=False,
... )
...
... def forward(self, input):
... pass
>>> class LlamaMLP(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.gate_up_proj = paddle.nn.Linear(
... model_config.hidden_size,
... model_config.intermediate_size * 2,
... bias_attr=False,
... )
...
... self.down_proj = paddle.nn.Linear(model_config.intermediate_size, model_config.hidden_size, bias_attr=False)
...
... def forward(self, input):
... pass
>>> class LlamaDecoderLayer(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.self_attn = LlamaAttention()
... self.mlp = LlamaMLP()
... self.input_layernorm = LlamaRMSNorm()
... self.post_attention_layernorm = LlamaRMSNorm()
...
... def forward(self, input):
... pass
>>> class LlamaModel(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.embedding = paddle.nn.Embedding(model_config.vocab_size, model_config.hidden_size)
... decoder_layers = []
... for _ in range(model_config.num_layers):
... decoder_layers.append(LlamaDecoderLayer())
...
... self.layers = paddle.nn.LayerList(decoder_layers)
... self.norm = LlamaRMSNorm()
...
... def forward(self, input):
... pass
>>> class LlamaLMHead(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.weight = self.create_parameter(
... shape=[model_config.hidden_size, model_config.vocab_size],
... dtype=paddle.get_default_dtype(),
... )
...
... def forward(self, input):
... pass
>>> class LlamaForCausalLM(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.llama = LlamaModel()
... self.lm_head = LlamaLMHead()
...
... def forward(self, input):
... pass
>>> mesh = dist.ProcessMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["dp", "mp", "pp"])
>>> dist.auto_parallel.set_mesh(mesh)
>>> parallel_config = {
... "dp_config": {'sharding_level': 1},
... "mp_config": {
... "parallelize_plan": {
... "llama.embed_tokens": [
... dist.ColWiseParallel(),
... dist.SequenceParallelBegin(),
... ],
... "llama.position_embedding": [
... dist.ColWiseParallel(),
... dist.SequenceParallelBegin(),
... ],
... "llama.layers.*.self_attn.qkv_proj": dist.ColWiseParallel(),
... "llama.layers.*.self_attn.o_proj": dist.RowWiseParallel(),
... "llama.layers.*.self_attn": dist.SequenceParallelDisable(),
... "llama.layers.*.mlp.gate_up_proj": dist.ColWiseParallel(),
... "llama.layers.*.mlp.down_proj": dist.RowWiseParallel(),
... "llama.layers.*.mlp": dist.SequenceParallelDisable(need_transpose=False),
... "lm_head.weight": dist.ColWiseParallel(),
... "lm_head": dist.SequenceParallelEnd(),
... }
... },
... "pp_config": {'split_spec': "llama.layers"},
... }
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
>>> model = LlamaForCausalLM()
>>> optimizer = paddle.optimizer.AdamW(parameters=model.parameters())
>>> dist_model, dist_optimizer = dist.parallelize(model, optimizer, config=parallel_config) # type: ignore[arg-type]
>>> # This case need to be executed in multi-card environment
>>> # python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 {test_case}.py
"""
if config is None:
warnings.warn(
"The `parallelize will do nothing since the config is `None`."
)
return model, optimizer
assert isinstance(config, dict)
if mesh is not None:
assert isinstance(mesh, core.ProcessMesh), (
"The mesh must be an instance of paddle.distributed.ProcessMesh."
)
g_mesh = fleet.auto.get_mesh()
if g_mesh is not None and g_mesh != mesh:
warnings.warn(
"The mesh set by `fleet.auto.set_mesh` is different with the mesh pass to "
"`parallelize`. Will overwrite the previous mesh"
)
fleet.auto.set_mesh(mesh)
pp_config = config.get('pp_config')
mp_config = config.get('mp_config')
dp_config = config.get('dp_config')
cp_config = config.get('cp_config')
if pp_config is not None:
assert isinstance(pp_config, dict)
model, optimizer = pipeline_parallel(
model,
optimizer,
pp_config,
)
if mp_config is not None:
assert isinstance(mp_config, dict)
if cp_config is not None:
assert isinstance(cp_config, dict)
assert "parallelize_plan" in cp_config.keys()
assert "parallelize_plan" in mp_config.keys()
mp_config['parallelize_plan'].update(cp_config['parallelize_plan'])
model, optimizer = tensor_parallel(model, optimizer, mp_config)
elif cp_config is not None:
assert isinstance(cp_config, dict)
model, optimizer = tensor_parallel(
model,
optimizer,
cp_config,
)
if dp_config is not None:
assert isinstance(dp_config, dict)
if 'sharding_level' not in dp_config.keys():
warnings.warn(
"The dp_config doesn't contain sharding_level, will run under dp."
)
model, optimizer = sharded_data_parallel(
model,
optimizer,
config=dp_config,
)
model, optimizer = parallelize_model_and_optimizer(model, optimizer)
return model, optimizer
has_parallelized_model = False
def parallelize_model(model, mesh=None, config=None):
if config is None:
warnings.warn(
"The `parallelize_model will do nothing since the config is `None`."
)
return model
assert isinstance(config, dict)
if mesh is not None:
assert isinstance(mesh, core.ProcessMesh), (
"The mesh must be an instance of paddle.distributed.ProcessMesh."
)
g_mesh = fleet.auto.get_mesh()
if g_mesh is not None and g_mesh != mesh:
warnings.warn(
"The mesh set by `fleet.auto.set_mesh` is different with the mesh pass to "
"`parallelize_model`. Will overwrite the previous mesh"
)
fleet.auto.set_mesh(mesh)
global has_parallelized_model
has_parallelized_model = True
model, _ = parallelize(model, None, mesh, config)
return model
def parallelize_optimizer(optimizer, mesh=None, config=None):
if config is None:
warnings.warn(
"The `parallelize_optimizer will do nothing since the config is `None`."
)
return optimizer
assert isinstance(config, dict)
if mesh is not None:
assert isinstance(mesh, core.ProcessMesh), (
"The mesh must be an instance of paddle.distributed.ProcessMesh."
)
g_mesh = fleet.auto.get_mesh()
if g_mesh is not None and g_mesh != mesh:
warnings.warn(
"The mesh set by `fleet.auto.set_mesh` is different with the mesh pass to "
"`parallelize_optimizer`. Will overwrite the previous mesh"
)
fleet.auto.set_mesh(mesh)
global has_parallelized_model
assert has_parallelized_model, (
"Please parallelize the model before parallelize optimizer."
)
param_list = optimizer._parameter_list
if isinstance(param_list[0], dict):
for param_group in param_list:
for param in param_group['params']:
assert param.is_dist(), (
"Please use model after parallelize to create optimizer."
)
else:
for param in param_list:
assert param.is_dist(), (
"Please use model after parallelize to create optimizer."
)
dp_config = config.get('dp_config')
level = None
sharding_mesh_dim = None
if dp_config is not None:
if 'sharding_level' not in dp_config.keys():
warnings.warn(
"The dp_config doesn't contain sharding_level, will run under dp."
)
level = dp_config.get('sharding_level')
sharding_mesh_dim = dp_config.get('sharding_mesh_dim', "dp")
optimizer = ParallelOptimizer(optimizer, level, sharding_mesh_dim)
optimizer = optimizer.parallelize()
return optimizer
python/paddle/distributed/auto_parallel/intermediate/pipeline_parallel.py
+419
View File
@@ -0,0 +1,419 @@
# 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
python/paddle/distributed/auto_parallel/intermediate/sharded_data_parallel.py
+88
View File
@@ -0,0 +1,88 @@
# 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.
from paddle.distributed import fleet
from .parallel_base import ParallelModel, ParallelOptimizer
class ShardedDataParallel(ParallelModel):
"""
ShardedDataParallel converts a single card model to a distributed data parallel model
Args:
model (paddle.nn.Layer): A single card model to be distributed.
optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed.
level (str): Zero stage, can be the following values:
0: no sharding (pure dp)
1: Zero Stage1
2: Zero Stage2
3: Zero Stage3
Default: None, which means optimizer is replicated among all process.
offload (bool): whether enable cpu offload strategy, not implemented currently.
exclude_layer (list): Specify which layers do not use the zero stage strategy, not implemented currently.
"""
def __init__(
self,
model,
offload=False,
exclude_layer=None,
):
super().__init__(model)
assert offload is False
assert exclude_layer is None
self.sharding_parallelizer = self.sharding_parallelizer_func
def sharding_parallelizer_func(self, model):
return model
def sharded_data_parallel(model, optimizer=None, config=None):
"""
sharded_data_parallel converts model and optimizer to distributed and supports set zero stage1/2/3
Args:
model (paddle.nn.Layer): A single card model to be distributed
optimizer (paddle.optimizer.Optimizer): an optimizer to be distributed
config (dict): {
"sharding_level": 0,
"offload": False,
"exclude_layer": None,
"sharding_mesh_dim": "dp",
}
Returns:
ShardedDataParallel: a distributed model
ParallelOptimizer: a distributed optimizer
"""
sdp_model = ShardedDataParallel(
model, bool(config.get('offload')), config.get('exclude_layer')
)
if optimizer is not None:
level = config.get('sharding_level')
sharding_mesh_dim = config.get('sharding_mesh_dim', "dp")
optimizer = ParallelOptimizer(optimizer, level, sharding_mesh_dim)
# check global_mesh
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 "dp" in mesh.dim_names, (
"dp must in the mesh dim_names when use sharded_data_parallel"
)
return sdp_model, optimizer
python/paddle/distributed/auto_parallel/intermediate/tensor_parallel.py
+954
View File
@@ -0,0 +1,954 @@
# 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.
from __future__ import annotations
import logging
import re
from typing import TYPE_CHECKING
import paddle
import paddle.distributed as dist
from .parallel_base import ParallelModel, ParallelOptimizer, is_tensor
if TYPE_CHECKING:
from collections.abc import Callable
from paddle import Tensor
from paddle.distributed import ProcessMesh
from paddle.nn import Layer
def c_split(x, process_mesh, need_transpose, split_type="sp"):
mp_index = process_mesh.dim_names.index('mp') # get the axis for the split
dp_index = process_mesh.dim_names.index('dp')
if isinstance(x, tuple):
target_x = x[0]
else:
target_x = x
assert is_tensor(target_x)
assert len(target_x.shape) == 3
if need_transpose:
target_x = paddle.transpose(target_x, perm=[1, 0, 2])
placements = target_x.placements
if placements is None:
placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
if split_type == "sp":
if placements[dp_index] == dist.Shard(0):
# NOTE(zhangwl):if shard(0) , input shape should be [b,s,h]
split_dims = dist.Shard(1)
elif placements[dp_index] == dist.Shard(1):
# NOTE(zhangwl):if shard(1) , input shape should be [s,b,h]
split_dims = dist.Shard(0)
else:
logging.warning(
f"parallel api don't know {target_x.shape} which dimension is batch, default is to cut to the 0th dimension"
)
split_dims = dist.Shard(0)
elif split_type == "mp":
split_dims = dist.Shard(2) # split h [b,s,h]
else:
raise ValueError(f"Unsupported split type {split_type}")
placements[mp_index] = split_dims
target_x = dist.reshard(target_x, process_mesh, placements)
if isinstance(x, tuple):
x = list(x)
x[0] = target_x
x = tuple(x)
else:
x = target_x
return x
def c_concat(x, process_mesh, need_transpose):
index = process_mesh.dim_names.index('mp') # get the axis for the split
if isinstance(x, tuple):
target_x = x[0]
else:
target_x = x
assert is_tensor(target_x)
assert len(target_x.shape) == 3
placements = target_x.placements
if placements is None:
placements = [dist.Replicate() for _ in range(len(process_mesh.shape))]
placements[index] = dist.Replicate()
target_x = dist.reshard(target_x, process_mesh, placements)
if need_transpose:
target_x = paddle.transpose(target_x, perm=[1, 0, 2])
if isinstance(x, tuple):
x = list(x)
x[0] = target_x
x = tuple(x)
else:
x = target_x
return x
class PlanBase:
def __init__(self):
self.share_param_list = {}
def apply(self, layer, process_mesh, shard_param_list):
raise NotImplementedError("Don't call the PlanBase directly.")
class ColWiseParallel(PlanBase):
"""
Col wise parallel plan for mp config.
Will try to split weight on the second dim and the bias on the first dim.
This api is designed for paddle.nn.Linear or paddle.nn.Embedding.
If any other instance of paddle.nn.Layer is passed,
this plan will try to split `layer.weight` and `layer.bias` if it has.
Note:
1. `layer.weight` should have two dims.
2. `layer.bias` should have one dim.
Args:
gather_output (bool): Whether gather the output to change it from a local tensor to a global tensor.
If gather the local tensor to global, an extra communication will be called.
The default value is `False`, which means keeping the output as a local tensor.
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()
>>> mp_config = {
... 'fc1': dist.ColWiseParallel(),
... }
"""
def __init__(self, gather_output: bool = False) -> None:
super().__init__()
self.gather_output = gather_output
def gather_output_hook(self, process_mesh):
def gather_hook(layer, input, output):
assert output is not None
return c_concat(output, process_mesh, False)
return gather_hook
def apply(self, layer, process_mesh, shard_param_list):
index = process_mesh.dim_names.index('mp') # get the axis for the split
size = len(process_mesh.shape)
placement = [dist.Replicate() for _ in range(size)]
param_placements = {}
assert isinstance(layer, paddle.nn.Layer)
if not isinstance(layer, (paddle.nn.Linear, paddle.nn.Embedding)):
logging.warning(
f"ColWiseParallel is designed to handle Linear and Embedding. "
f"But got {layer.__class__.__name__}. "
f"Will try to shard weight and bias if the layer contains one."
)
shard_param_list = set(shard_param_list)
if len(shard_param_list) == 0:
shard_param_list.add("weight")
shard_param_list.add("bias")
def shard_param(param_name):
if (
hasattr(layer, param_name)
and getattr(layer, param_name) is not None
):
layer_param = getattr(layer, param_name)
if layer_param.is_dist():
return
if len(layer_param.shape) == 2:
placement[index] = dist.Shard(1)
elif len(layer_param.shape) == 1:
placement[index] = dist.Shard(0)
else:
raise ValueError(f"{layer_param} should have 1 or 2 dims.")
# NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
if (
self.share_param_list is not None
and layer_param.name in self.share_param_list
and self.share_param_list[layer_param.name] > 1
):
param_placements.update({param_name: placement})
else:
layer_param = dist.shard_tensor(
layer_param,
process_mesh,
placement,
)
setattr(layer, param_name, layer_param)
for param_name in shard_param_list:
shard_param(param_name)
if self.gather_output:
layer.register_forward_post_hook(
self.gather_output_hook(process_mesh)
)
return param_placements
class RowWiseParallel(PlanBase):
"""
Row wise parallel plan for mp config.
Will try to split weight on the first dim.
This api is designed for paddle.nn.Linear or paddle.nn.Embedding.
If any other instance of paddle.nn.Layer is passed, this plan will try to split `layer.weight` if it has.
Note:
`layer.weight` should have two dims.
Args:
is_input_parallel (bool): Whether the input is a local tensor or a global tensor. If the input is a
global tensor, an extra split will be called. The default value is `True`,
which means the input is a local tensor.
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()
>>> mp_config = {
... 'fc1': dist.RowWiseParallel(),
... }
"""
def __init__(self, is_input_parallel: bool = True) -> None:
super().__init__()
self.is_input_parallel = is_input_parallel
def split_input_hook(self, process_mesh):
def split_hook(layer, input):
return c_split(input, process_mesh, False, split_type="mp")
return split_hook
def apply(self, layer, process_mesh, shard_param_list):
index = process_mesh.dim_names.index('mp') # get the axis for the split
size = len(process_mesh.shape)
placement = [dist.Replicate() for _ in range(size)]
placement[index] = dist.Shard(0)
param_placements = {}
assert isinstance(layer, paddle.nn.Layer)
if not isinstance(layer, (paddle.nn.Linear, paddle.nn.Embedding)):
logging.warning(
f"RowWiseParallel is designed to handle Linear and Embedding. "
f"But got {layer.__class__.__name__}. "
f"Will try to shard weight if the layer contains one."
)
shard_param_list = set(shard_param_list)
shard_param_list.discard("bias")
if len(shard_param_list) == 0:
shard_param_list.add("weight")
def shard_param(param_name):
if (
hasattr(layer, param_name)
and getattr(layer, param_name) is not None
):
layer_param = getattr(layer, param_name)
if layer_param.is_dist():
return
if len(layer_param.shape) != 2:
raise ValueError(f"{layer_param} should have 2 dims.")
# NOTE(zhangweilong):for share parameter, the parameter should be handled uniformly in the end
if (
self.share_param_list is not None
and layer_param.name in self.share_param_list
and self.share_param_list[layer_param.name] > 1
):
param_placements.update({param_name: placement})
else:
layer_param = dist.shard_tensor(
layer_param,
process_mesh,
placement,
)
setattr(layer, param_name, layer_param)
for param_name in shard_param_list:
shard_param(param_name)
if not self.is_input_parallel:
layer.register_forward_pre_hook(self.split_input_hook(process_mesh))
return param_placements
class PrepareLayerInput(PlanBase):
"""
Prepare the input of specific layer. User should provide one callable function.
Args:
fn (callable): A function that prepare the layer input. The function should take exactly
one parameter named `process_mesh` and return the pre hook.
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))
>>> def layer_input_hook(process_mesh):
... def hook(layer, input, output):
... return input
...
... return hook
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
>>> layer = MLP()
>>> mp_config = {
... 'fc1': dist.PrepareLayerOutput(layer_input_hook),
... }
"""
def __init__(
self,
fn: (
Callable[
[ProcessMesh],
Callable[
[Layer, tuple[Tensor], tuple[Tensor]], [tuple[Tensor]]
],
]
| None
) = None,
) -> None:
super().__init__()
assert callable(fn)
self.fn = fn
def apply(self, layer, process_mesh, shard_param_list):
layer.register_forward_pre_hook(self.fn(process_mesh=process_mesh))
class PrepareLayerOutput(PlanBase):
"""
Prepare the output of specific layer. User should provide one callable function.
Args:
fn (callable): A function that prepare the layer input. The function should take exactly
one parameter named `process_mesh` and return the post hook.
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))
>>> def layer_output_hook(process_mesh):
... def hook(layer, input, output):
... return output
...
... return hook
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
>>> layer = MLP()
>>> mp_config = {
... 'fc1': dist.PrepareLayerOutput(layer_output_hook),
... }
"""
def __init__(
self,
fn: (
Callable[
[ProcessMesh],
Callable[
[Layer, tuple[Tensor], tuple[Tensor]], [tuple[Tensor]]
],
]
| None
) = None,
) -> None:
super().__init__()
assert callable(fn)
self.fn = fn
def apply(self, layer, process_mesh, shard_param_list):
layer.register_forward_post_hook(self.fn(process_mesh=process_mesh))
class SequenceParallelBegin(PlanBase):
"""
Sequence parallel plan for mp config.
This plan marks the beginning of the sp and should be added to the LAST layer before the sp range.
Note:
DON'T mark any layer in the sp range.
Args:
need_transpose (bool): the default value is `True`. With `need_transpose=True`, this plan will transfer
the output from [b, s, h] to [s/mp, b, h]. With `need_transpose=False`, this plan will transfer
the output from [s, b, h] to [s/mp, b, h].
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()
>>> mp_config = {
... 'fc1': dist.SequenceParallelBegin(),
... }
"""
def __init__(self, need_transpose: bool = True) -> None:
super().__init__()
self.need_transpose = need_transpose
def sequence_parallel_begin(self, process_mesh):
def begin(layer, input, output):
assert output is not None
return c_split(output, process_mesh, self.need_transpose)
return begin
def apply(self, layer, process_mesh, shard_param_list):
layer.register_forward_post_hook(
self.sequence_parallel_begin(process_mesh)
)
class SequenceParallelEnd(PlanBase):
"""
Sequence parallel plan for mp config.
This plan marks the ending of the sp and should be added to the FIRST layer after the sp range.
Note:
DON'T mark any layer in the sp range.
Args:
need_transpose (bool): the default value is `True`. With `need_transpose=True`, this plan will transfer
the input from [s/mp, b, h] to [b, s, h]. With `need_transpose=False`, this plan will transfer the
input from [s/mp, b, h] to [s, b, h].
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()
>>> mp_config = {
... 'fc1': dist.SequenceParallelEnd(),
... }
"""
def __init__(self, need_transpose: bool = True) -> None:
super().__init__()
self.need_transpose = need_transpose
def sequence_parallel_end(self, process_mesh):
def end(layer, input, output=None):
assert input is not None
return c_concat(input, process_mesh, self.need_transpose)
return end
def apply(self, layer, process_mesh, shard_param_list):
layer.register_forward_pre_hook(
self.sequence_parallel_end(process_mesh)
)
class SequenceParallelEnable(PlanBase):
"""
Sequence parallel plan for mp config.
Do sequence parallel on the layer. Note the input should be in [b, s, h] format.
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()
>>> mp_config = {
... 'fc1': dist.SequenceParallelEnable(),
... }
"""
def __init__(self) -> None:
super().__init__()
def sequence_parallel_begin(self, process_mesh):
def begin(layer, input, output=None):
assert input is not None
return c_split(input, process_mesh, True)
return begin
def sequence_parallel_end(self, process_mesh):
def end(layer, input, output):
assert output is not None
return c_concat(output, process_mesh, True)
return end
def apply(self, layer, process_mesh, shard_param_list):
logging.warning(
"Sequence parallel with the usage of SequenceParallel may not reach the best throughput. "
"Try to use SequenceParallelBegin/End to achieve better performance"
)
layer.register_forward_pre_hook(
self.sequence_parallel_begin(process_mesh)
)
layer.register_forward_post_hook(
self.sequence_parallel_end(process_mesh)
)
class SequenceParallelDisable(PlanBase):
"""
Sequence parallel plan for mp config.
Disable sequence parallel on the layer.
Args:
need_transpose (bool): the default value is `True`. If the need_transpose is `True`: this plan will transfer
the input from [s/mp, b, h] to [b, s, h] and then transfer the output from [b, s, h] to [s/mp, b, h].
If the need_transpose is `False`: this plan will transfer the input from [s/mp, b, h] to [s, b, h] and
then transfer the output from [s, b, h] to [s/mp, b, h].
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()
>>> mp_config = {
... 'fc1': dist.SequenceParallelDisable(),
... }
"""
def __init__(self, need_transpose: bool = True) -> None:
super().__init__()
self.need_transpose = need_transpose
def sequence_parallel_begin(self, process_mesh):
def begin(layer, input, output=None):
return c_split(output, process_mesh, self.need_transpose)
return begin
def sequence_parallel_end(self, process_mesh):
def end(layer, input, output=None):
return c_concat(input, process_mesh, self.need_transpose)
return end
def apply(self, layer, process_mesh, shard_param_list):
layer.register_forward_pre_hook(
self.sequence_parallel_end(process_mesh)
)
layer.register_forward_post_hook(
self.sequence_parallel_begin(process_mesh)
)
class ConvParallel(PlanBase):
"""
A strategy for enabling spatial parallelism on ``paddle.nn.Conv2D`` layers
by sharding the input tensor along its Width (W) dimension.
When this ``ConvParallel`` configuration is applied to a ``Conv2D`` layer,
the layer's input tensor will have its width dimension split across devices
in the model parallel group. This can help reduce memory usage from activations,
especially when dealing with inputs that have a large width.
To enable width-wise input sharding correctly, make sure your `Conv2D` layer
satisfies the following conditions along the width dimension:
- **Dilation** must be set to `1`.
- **If no width padding is used:**
- The input width must be evenly divisible by the stride width.
- The stride width must be equal to the kernel width.
- **If width padding is used:**
- The stride width must be `1`.
- The total input width must be at least half the kernel width.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn as nn
>>> import paddle.distributed as dist
>>> class SimpleConvNet(nn.Layer):
... def __init__(self, data_format="NCHW"):
... super().__init__()
... self.conv1 = nn.Conv2D(
... 3,
... 8,
... kernel_size=3,
... padding=1,
... data_format=data_format,
... )
... self.relu = nn.ReLU()
...
... def forward(self, x):
... x = self.conv1(x)
... return self.relu(x)
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
>>> model = SimpleConvNet(data_format="NCHW")
>>> mp_config = {
... "parallelize_plan": {
... "conv1": dist.ConvParallel(),
... },
... }
"""
def __init__(self) -> None:
super().__init__()
@staticmethod
def _is_supported(
input_size,
kernel_size,
stride,
padding,
dilation,
data_format,
mp_group_size,
):
idx_w_input = -1
idx_w_kernel = -1
if data_format == "NCHW":
idx_w_input = 3
idx_w_kernel = 3
elif data_format == "NHWC":
idx_w_input = 2
idx_w_kernel = 3
else:
return False
if input_size[idx_w_input] % mp_group_size != 0:
return False
dilation_w = dilation[1]
padding_w = padding[1]
stride_w = stride[1]
input_w = input_size[idx_w_input]
kernel_w = kernel_size[idx_w_kernel]
if dilation_w != 1:
# RingConv2d only supports dilation=1.
# Larger dilation would require enlarged halo regions and more complex communication.
return False
if padding_w == 0:
# To avoid halo exchange when padding=0, we require:
# - input_w must be divisible by stride_w, so partitions align evenly across ranks.
# - stride_w == kernel_w, so each kernel operates on disjoint local regions.
if input_w % stride_w != 0:
return False
if stride_w != kernel_w:
return False
else:
# When padding > 0, halo exchange is needed.
# To simplify halo logic, we require:
# - stride_w == 1: ensures each output element is computed from overlapping input,
# and no input region is skipped, simplifying halo construction.
# - kernel_w // 2 <= input_w: prevents the kernel from exceeding local input.
if stride_w != 1:
return False
if kernel_w // 2 > input_w:
return False
return True
def conv_parallel_start(self, process_mesh, data_format):
def start(layer, input, output=None):
if data_format == "NCHW":
shard_w_dim = 3
elif data_format == "NHWC":
shard_w_dim = 2
else:
raise ValueError(
f"Unsupported data_format: {data_format}. "
"Only NCHW and NHWC are supported."
)
if isinstance(input, tuple):
x = input[0]
else:
x = input
placements = x.placements
mp_index = process_mesh.dim_names.index('mp')
mp_group_size = process_mesh.get_dim_size('mp')
# Note(luchang): for intermediate api, when this ConvLayer is
# not supported, we just skip apply parallelization.
if not ConvParallel._is_supported(
x.shape,
layer.weight.shape,
layer._stride,
layer._updated_padding,
layer._dilation,
data_format,
mp_group_size,
):
return input
if placements is None:
placements = [
dist.Replicate() for _ in range(len(process_mesh.shape))
]
if placements[mp_index] == dist.Shard(shard_w_dim):
return input
placements[mp_index] = dist.Shard(shard_w_dim)
if not x.is_dist():
x = dist.shard_tensor(x, process_mesh, placements)
else:
x = dist.reshard(x, process_mesh, placements)
if isinstance(input, tuple):
input = list(input)
input[0] = x
input = tuple(input)
else:
input = x
return input
return start
def apply(self, layer, process_mesh, shard_param_list):
layer.register_forward_pre_hook(
self.conv_parallel_start(process_mesh, layer._data_format)
)
class TensorParallel(ParallelModel):
def __init__(self, model, parallelize_plan=None):
super().__init__(model)
if parallelize_plan is not None:
assert isinstance(parallelize_plan, dict)
for key, plan in parallelize_plan.items():
assert isinstance(key, str), (
"The key of the parallelize plan should be a string."
)
if not isinstance(plan, list):
plan = [plan]
for p in plan:
assert isinstance(p, PlanBase), (
"The value the the parallelize plan should be a instance of PlanBase or a list of PlanBase."
)
self.global_mesh = dist.auto_parallel.get_mesh()
self.parallelize_plan = parallelize_plan
self.tp_parallelizer = self.tensor_parallelizer_fn
def match_layer(self, layer, name):
# Match the layer to a plan.
# Will return the plan if the layer hits one, otherwise return None.
plans = []
for key, plan in self.parallelize_plan.items():
attr_name = key.split('.')[-1]
shard_param_list = []
# Find some plan for specific parameter, such as
# "lm_head.weight": ColWiseParallel()
# "qkv_proj.lora_A" ColWiseParallel()
# if there is no plan for specific parameter, layer will be sharded by default: layer.weight and layer.bias
if key.endswith(f".{attr_name}"):
if hasattr(layer, attr_name) and is_tensor(
getattr(layer, attr_name)
):
key = key.replace(f".{attr_name}", "")
shard_param_list.append(attr_name)
re_find = re.match(key, name)
if key == name or (
re_find is not None
and int(re_find.end()) - int(re_find.start()) == len(name)
):
if isinstance(plan, PlanBase):
plan = [plan]
plans.append([plan, shard_param_list])
return plans
def tensor_parallelizer_fn(self, model):
if self.parallelize_plan is None:
return
layer_param_placements = {}
share_param_list = {}
for name, layer in model.named_sublayers():
for param_name in list(layer._parameters.keys()):
param = getattr(layer, param_name)
if param.name not in share_param_list:
share_param_list[param.name] = 1
continue
share_param_list[param.name] += 1
for name, layer in model.named_sublayers():
plans = self.match_layer(layer, name)
layer_param_placements[layer] = {}
if len(plans) > 0:
pp_idx = getattr(layer, "pipeline_stage_index", 0)
for plan in plans:
real_plan, shard_param_list = plan
for p in real_plan:
p.share_param_list = share_param_list
param_placements = p.apply(
layer, self.get_mesh(pp_idx), shard_param_list
)
if param_placements is not None and param_placements:
layer_param_placements[layer].update(
param_placements
)
return model, layer_param_placements
def tensor_parallel(model, optimizer=None, config=None):
"""
Tensor parallel.
Args:
model (paddle.nn.Layer): the model to be shard into tensor parallel.
optimizer (paddle.optimizer.Optimizer): the optimizer.
config (dict): {
"parallelize_plan": dict, the plan to shard the layer.
}
Returns:
model: model after tp
optimizer: optimizer after tp
NOTE: the plan should be a dict maps layer name or parameter name to a split_plan,
which will be used to split the layer or the parameter. The name can be written in regular format.
An example for the plan is:
```
plan = {
"llama.embed_tokens": ColWiseParallel(),
"llama.layers.*.self_attn.q_proj": ColWiseParallel(),
"llama.layers.*.self_attn.k_proj": ColWiseParallel(),
"llama.layers.*.self_attn.v_proj": ColWiseParallel(),
"llama.layers.*.self_attn.o_proj": RowWiseParallel(),
"llama.layers.*.mlp.gate_proj": ColWiseParallel(),
"llama.layers.*.mlp.up_proj": ColWiseParallel(),
"llama.layers.*.mlp.down_proj": RowWiseParallel(),
"lm_head.weight": ColWiseParallel(),
}
```
"""
parallelize_plan = config.get("parallelize_plan")
if parallelize_plan is None:
# Do nothing if no plan.
logging.warning(
"No parallelize plan, tensor parallel won't do anything."
)
return model, optimizer
global_mesh = dist.auto_parallel.get_mesh()
assert global_mesh is not None, (
"global mesh must not be None, please call fleet.auto.set_mesh(global_mesh) firstly"
)
assert "mp" in global_mesh.dim_names, (
"mp must in the mesh dim_names when use tensor_parallel"
)
model = TensorParallel(model, parallelize_plan)
if optimizer is not None:
optimizer = ParallelOptimizer(optimizer)
return model, optimizer