chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user