chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) 2022 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 .mp_layers import ( # noqa: F401
|
||||
ColumnParallelLinear,
|
||||
ParallelCrossEntropy,
|
||||
RowParallelLinear,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from .random import ( # noqa: F401
|
||||
RNGStatesTracker,
|
||||
dropout,
|
||||
get_rng_state_tracker,
|
||||
model_parallel_random_seed,
|
||||
)
|
||||
@@ -0,0 +1,886 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
|
||||
import paddle
|
||||
from paddle.autograd import PyLayer
|
||||
from paddle.base import core
|
||||
from paddle.distributed import fleet
|
||||
from paddle.nn import functional as F
|
||||
|
||||
from ....communication.reduce import ReduceOp, _get_reduce_op
|
||||
from ....flex_checkpoint.dcp.sharded_weight import build_sharded_state_dict
|
||||
from ...base import topology as tp
|
||||
from ...utils.log_util import logger
|
||||
from . import mp_ops
|
||||
from .random import get_rng_state_tracker
|
||||
|
||||
__all__ = []
|
||||
|
||||
# Follow this paper to achieve the file:
|
||||
# Shoeybi M, Patwary M, Puri R, et al. Megatron-lm: Training multi-billion parameter
|
||||
# language models using model parallelism[J]. arXiv preprint arXiv:1909.08053, 2019. (https://arxiv.org/abs/1909.08053)
|
||||
|
||||
|
||||
def is_fused_matmul_bias_supported():
|
||||
return hasattr(core.eager.ops.legacy, 'fused_gemm_epilogue')
|
||||
|
||||
|
||||
def is_fused_linear_param_grad_add_supported():
|
||||
if (
|
||||
paddle.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm()
|
||||
) or paddle.is_compiled_with_xpu():
|
||||
return hasattr(paddle._C_ops, 'fused_linear_param_grad_add')
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
class VocabParallelEmbedding(paddle.nn.Layer):
|
||||
"""Embedding mp parallelized in the vocabulary dimension.
|
||||
this class is used for splitting embedding in mp group.
|
||||
|
||||
Args:
|
||||
num_embeddings(int): One element which indicate the size of the dictionary of embeddings.
|
||||
embedding_dim(int): One element which indicate the size of each embedding vector respectively.
|
||||
weight_attr(ParamAttr|None): To specify the weight parameter property. Default: None, which means the
|
||||
default weight parameter property is used. See usage for details in :ref:`api_paddle_ParamAttr` . In addition,
|
||||
user-defined or pre-trained word vectors can be loaded with the :attr:`param_attr` parameter.
|
||||
The local word vector needs to be transformed into numpy format, and the shape of local word
|
||||
vector should be consistent with :attr:`num_embeddings` . Then :ref:`api_paddle_nn_initializer_Assign`
|
||||
is used to load custom or pre-trained word vectors. See code example for details.
|
||||
mp_group(Group): The tensor parallel group.
|
||||
name(str, optional): For detailed information, please refer
|
||||
to :ref:`api_guide_Name`. Usually name is no need to set and
|
||||
None by default.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed import fleet
|
||||
|
||||
>>> class SimpleMPNet(paddle.nn.Layer):
|
||||
... def __init__(self, vocab_size, hidden_size, inner_size, output_size):
|
||||
... super().__init__()
|
||||
... self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
|
||||
... hidden_size,
|
||||
... inner_size,
|
||||
... gather_output=False,
|
||||
... has_bias=True,
|
||||
... )
|
||||
... self.linear2 = fleet.meta_parallel.RowParallelLinear(
|
||||
... inner_size,
|
||||
... hidden_size,
|
||||
... input_is_parallel=True,
|
||||
... has_bias=True,
|
||||
... )
|
||||
... self.linear3 = paddle.nn.Linear(hidden_size, output_size)
|
||||
... self.embedding = fleet.meta_parallel.VocabParallelEmbedding(vocab_size, hidden_size)
|
||||
...
|
||||
... def forward(self, x):
|
||||
... x = self.embedding(x)
|
||||
... x = self.linear1(x)
|
||||
... x = self.linear2(x)
|
||||
... x = self.linear3(x)
|
||||
... return x
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
weight_attr=None,
|
||||
mp_group=None,
|
||||
name=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.model_parallel_group = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
|
||||
if mp_group is None
|
||||
else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
self.rank = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
|
||||
if mp_group is None
|
||||
else mp_group.rank
|
||||
)
|
||||
|
||||
self.origin_num_embeddings = num_embeddings
|
||||
self.is_mp = self.world_size > 1
|
||||
|
||||
assert num_embeddings % self.world_size == 0, (
|
||||
"The length of the vocabulary must be divisible by the parallelism degree of MP"
|
||||
)
|
||||
|
||||
per_part_size = num_embeddings // self.world_size
|
||||
|
||||
self.vocab_start_index = self.rank * per_part_size
|
||||
self._dtype = self._helper.get_default_dtype()
|
||||
self._size = [per_part_size, embedding_dim]
|
||||
self._weight_attr = weight_attr
|
||||
self._name = name
|
||||
self.num_embeddings = num_embeddings
|
||||
|
||||
if self.is_mp and paddle.in_dynamic_mode():
|
||||
with get_rng_state_tracker().rng_state():
|
||||
self.weight = self.create_parameter(
|
||||
attr=self._weight_attr,
|
||||
shape=self._size,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
else:
|
||||
self.weight = self.create_parameter(
|
||||
attr=self._weight_attr,
|
||||
shape=self._size,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
|
||||
self.weight.is_distributed = True if self.is_mp else False
|
||||
if self.weight.is_distributed:
|
||||
self.weight.split_axis = 0
|
||||
|
||||
def forward(self, x):
|
||||
if self.is_mp:
|
||||
output_parallel = mp_ops._c_lookup_table(
|
||||
self.weight,
|
||||
x,
|
||||
start_index=self.vocab_start_index,
|
||||
vocab_size=self.num_embeddings,
|
||||
name=self._name,
|
||||
)
|
||||
output = mp_ops._mp_allreduce(
|
||||
output_parallel,
|
||||
group=self.model_parallel_group,
|
||||
use_calc_stream=True,
|
||||
use_model_parallel=True,
|
||||
)
|
||||
else:
|
||||
output = F.embedding(
|
||||
x,
|
||||
weight=self.weight,
|
||||
padding_idx=None,
|
||||
sparse=False,
|
||||
name=self._name,
|
||||
)
|
||||
return output
|
||||
|
||||
def sharded_state_dict(
|
||||
self,
|
||||
structured_name_prefix: str = "",
|
||||
):
|
||||
state_dict = self.state_dict(structured_name_prefix="")
|
||||
return build_sharded_state_dict(
|
||||
state_dict, {"weight": 0}, structured_name_prefix
|
||||
)
|
||||
|
||||
|
||||
_raise_cuda_env_unset_warning = True
|
||||
|
||||
|
||||
class InnerOverlapLinear(paddle.autograd.PyLayer):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
fuse_matmul_bias,
|
||||
mp_async_allreduce,
|
||||
mp_skip_c_identity,
|
||||
mp_fused_linear_param_grad_add,
|
||||
model_parallel_group,
|
||||
):
|
||||
ctx.save_for_backward(x, weight, bias)
|
||||
ctx.model_parallel_group = model_parallel_group
|
||||
ctx.mp_fused_linear_param_grad_add = mp_fused_linear_param_grad_add
|
||||
if mp_skip_c_identity is False:
|
||||
x = paddle._legacy_C_ops.c_identity(
|
||||
x,
|
||||
'use_calc_stream',
|
||||
True,
|
||||
'ring_id',
|
||||
model_parallel_group.id,
|
||||
'use_model_parallel',
|
||||
True,
|
||||
)
|
||||
if not fuse_matmul_bias:
|
||||
return paddle._C_ops.linear(x, weight, bias)
|
||||
else:
|
||||
return paddle._legacy_C_ops.fused_gemm_epilogue(x, weight, bias)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy):
|
||||
x, weight, bias = ctx.saved_tensor()
|
||||
if dy.dtype == weight.dtype:
|
||||
dx = paddle.matmul(dy, weight, transpose_y=True)
|
||||
else:
|
||||
dx = paddle.matmul(
|
||||
dy, paddle.cast(weight, dtype=dy.dtype), transpose_y=True
|
||||
)
|
||||
op_type = _get_reduce_op(ReduceOp.SUM)
|
||||
task = ctx.model_parallel_group.process_group.all_reduce(
|
||||
dx, op_type, sync_op=False
|
||||
)
|
||||
# Using small operation to preempt GPU SMs for all_reduce to achieve overlap.
|
||||
if int(os.getenv("CUDA_DEVICE_MAX_CONNECTIONS", "0")) != 1:
|
||||
global _raise_cuda_env_unset_warning
|
||||
if _raise_cuda_env_unset_warning:
|
||||
logger.warning(
|
||||
"You set mp_async_allreduce=True, but you forget to set environment "
|
||||
"variable CUDA_DEVICE_MAX_CONNECTIONS=1, which may leads to performance "
|
||||
"loss. Try to export CUDA_DEVICE_MAX_CONNECTIONS=1 for better performance."
|
||||
)
|
||||
_raise_cuda_env_unset_warning = False
|
||||
tmp = paddle.ones([512])
|
||||
|
||||
if ctx.mp_fused_linear_param_grad_add:
|
||||
if not is_fused_linear_param_grad_add_supported():
|
||||
raise NotImplementedError(
|
||||
"You set mp_fused_linear_param_grad_add=True, "
|
||||
"however, the paddle you are using not support this operation. "
|
||||
"Please unset fused_linear_param_grad_add or use paddle compiled "
|
||||
"with cuda 11.6 or higher."
|
||||
)
|
||||
|
||||
if bias is None:
|
||||
if hasattr(weight, "main_grad"):
|
||||
(
|
||||
weight.main_grad,
|
||||
_,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
x, dy, weight.main_grad, None, True, False
|
||||
)
|
||||
task.wait()
|
||||
return dx, None
|
||||
else:
|
||||
if weight.grad is not None:
|
||||
(
|
||||
weight.grad,
|
||||
_,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
x, dy, weight.grad, None, False, False
|
||||
)
|
||||
task.wait()
|
||||
return dx, None
|
||||
else:
|
||||
(
|
||||
dw,
|
||||
_,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
x, dy, None, None, False, False
|
||||
)
|
||||
task.wait()
|
||||
return dx, dw
|
||||
|
||||
if hasattr(weight, "main_grad") and hasattr(bias, "main_grad"):
|
||||
(
|
||||
weight.main_grad,
|
||||
bias.main_grad,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
x,
|
||||
dy,
|
||||
weight.main_grad,
|
||||
bias.main_grad,
|
||||
True,
|
||||
True,
|
||||
)
|
||||
task.wait()
|
||||
return dx, None, None
|
||||
else:
|
||||
if weight.grad is not None:
|
||||
assert bias.grad is not None
|
||||
(
|
||||
weight.grad,
|
||||
bias.grad,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
x, dy, weight.grad, bias.grad, False, True
|
||||
)
|
||||
task.wait()
|
||||
return dx, None, None
|
||||
else:
|
||||
# When main_grad is not enabled and gradient_accumulation is used, the grad is not initialized for the first acc step.
|
||||
(
|
||||
dw,
|
||||
dbias,
|
||||
) = paddle._C_ops.fused_linear_param_grad_add(
|
||||
x, dy, None, None, False, True
|
||||
)
|
||||
task.wait()
|
||||
return dx, dw, dbias
|
||||
else:
|
||||
dy = dy.reshape([-1, dy.shape[-1]])
|
||||
dw = paddle.matmul(
|
||||
x.reshape([-1, x.shape[-1]]),
|
||||
dy,
|
||||
transpose_x=True,
|
||||
)
|
||||
if bias is None:
|
||||
task.wait()
|
||||
return dx, dw
|
||||
else:
|
||||
dbias = paddle.sum(dy, axis=0)
|
||||
task.wait()
|
||||
return dx, dw, dbias
|
||||
|
||||
|
||||
class ColumnParallelLinear(paddle.nn.Layer):
|
||||
"""Linear layer with mp parallelized(column).
|
||||
this class is used for splitting Linear Layer in mp group, column split the weight of the Linear layer.
|
||||
|
||||
Args:
|
||||
in_features(int): The number of input units.
|
||||
out_features(int): The number of output units.
|
||||
weight_attr(ParamAttr|None): The attribute for the learnable weight of this layer. The default value is None
|
||||
and the weight will be initialized to zero. For detailed information, please refer to paddle.ParamAttr.
|
||||
has_bias(bool): whether to add bias.
|
||||
gather_output(bool): whether to do allgather for the output of each rank.
|
||||
fuse_matmul_bias(bool): whether to fuse matmul and bias.
|
||||
mp_group(Group): The tensor parallel group.
|
||||
name(str, optional): Normally there is no need for user to set this parameter.
|
||||
For detailed information, please refer to :ref:`api_guide_Name` .
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed import fleet
|
||||
|
||||
>>> class SimpleMPNet(paddle.nn.Layer):
|
||||
... def __init__(self, vocab_size, hidden_size, inner_size, output_size):
|
||||
... super().__init__()
|
||||
... self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
|
||||
... hidden_size,
|
||||
... inner_size,
|
||||
... gather_output=False,
|
||||
... has_bias=True,
|
||||
... )
|
||||
... self.linear2 = fleet.meta_parallel.RowParallelLinear(
|
||||
... inner_size,
|
||||
... hidden_size,
|
||||
... input_is_parallel=True,
|
||||
... has_bias=True,
|
||||
... )
|
||||
... self.linear3 = paddle.nn.Linear(hidden_size, output_size)
|
||||
... self.embedding = fleet.meta_parallel.VocabParallelEmbedding(vocab_size, hidden_size)
|
||||
...
|
||||
... def forward(self, x):
|
||||
... x = self.embedding(x)
|
||||
... x = self.linear1(x)
|
||||
... x = self.linear2(x)
|
||||
... x = self.linear3(x)
|
||||
... return x
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
weight_attr=None,
|
||||
has_bias=None,
|
||||
gather_output=True,
|
||||
fuse_matmul_bias=False,
|
||||
mp_group=None,
|
||||
name=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.model_parallel_group = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
|
||||
if mp_group is None
|
||||
else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
self._name = name
|
||||
self.is_mp = self.world_size > 1
|
||||
|
||||
self.gather_output = gather_output
|
||||
assert out_features % self.world_size == 0, (
|
||||
f"Number of column of the weight for linear ({out_features}) must be"
|
||||
f" divisible by model parallel size ({self.world_size})"
|
||||
)
|
||||
self.output_size_per_partition = out_features // self.world_size
|
||||
|
||||
self._weight_attr = weight_attr
|
||||
self._dtype = self._helper.get_default_dtype()
|
||||
|
||||
if self.is_mp and paddle.in_dynamic_mode():
|
||||
with get_rng_state_tracker().rng_state():
|
||||
self.weight = self.create_parameter(
|
||||
shape=[in_features, self.output_size_per_partition],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
else:
|
||||
self.weight = self.create_parameter(
|
||||
shape=[in_features, self.output_size_per_partition],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
|
||||
self.weight.is_distributed = True if self.is_mp else False
|
||||
|
||||
if self.weight.is_distributed:
|
||||
self.weight.split_axis = 1
|
||||
|
||||
if has_bias:
|
||||
# initialize bias to zero like Megatron
|
||||
self.bias = self.create_parameter(
|
||||
shape=[self.output_size_per_partition],
|
||||
attr=paddle.nn.initializer.Constant(value=0.0),
|
||||
dtype=self._dtype,
|
||||
is_bias=True,
|
||||
)
|
||||
self.bias.is_distributed = True if self.is_mp else False
|
||||
if self.bias.is_distributed:
|
||||
self.bias.split_axis = 0
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
self.linear = F.linear
|
||||
|
||||
self.fuse_matmul_bias = fuse_matmul_bias
|
||||
|
||||
mp_configs = fleet.fleet._user_defined_strategy.hybrid_configs[
|
||||
"mp_configs"
|
||||
]
|
||||
self.mp_async_allreduce = self.is_mp and mp_configs.mp_async_allreduce
|
||||
self.mp_skip_c_identity = (
|
||||
self.is_mp
|
||||
and mp_configs.mp_async_allreduce
|
||||
and mp_configs.mp_skip_c_identity
|
||||
)
|
||||
self.mp_fused_linear_param_grad_add = (
|
||||
self.is_mp
|
||||
and mp_configs.mp_async_allreduce
|
||||
and mp_configs.mp_fused_linear_param_grad_add
|
||||
)
|
||||
if (
|
||||
self.mp_async_allreduce
|
||||
or self.mp_skip_c_identity
|
||||
or self.mp_fused_linear_param_grad_add
|
||||
):
|
||||
assert paddle.in_dynamic_mode(), (
|
||||
"mp_async_allreduce, mp_skip_c_identity and mp_fused_linear_param_grad_add are only available under dygraph mode"
|
||||
)
|
||||
if self.fuse_matmul_bias:
|
||||
if not is_fused_matmul_bias_supported():
|
||||
raise NotImplementedError(
|
||||
"You set fuse_matmul_bias=True in ColumnParallelLinear, "
|
||||
"however, the paddle you are using not support this operation. "
|
||||
"Please set fuse_matmul_bias=False or use paddle compiled "
|
||||
"with cuda 11.6 or higher."
|
||||
)
|
||||
from paddle.incubate.nn.functional import fused_linear
|
||||
|
||||
self.linear = fused_linear
|
||||
|
||||
def forward(self, x):
|
||||
# use inner api to process identity
|
||||
|
||||
def _overlap_linear():
|
||||
return InnerOverlapLinear.apply(
|
||||
x,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.fuse_matmul_bias,
|
||||
self.mp_async_allreduce,
|
||||
self.mp_skip_c_identity,
|
||||
self.mp_fused_linear_param_grad_add,
|
||||
self.model_parallel_group,
|
||||
)
|
||||
|
||||
if self.mp_async_allreduce:
|
||||
output_parallel = _overlap_linear()
|
||||
else:
|
||||
if self.is_mp:
|
||||
input_parallel = mp_ops._c_identity(
|
||||
x,
|
||||
group=self.model_parallel_group,
|
||||
skip_c_identity_dynamic=self.mp_skip_c_identity,
|
||||
)
|
||||
else:
|
||||
input_parallel = x
|
||||
|
||||
output_parallel = self.linear(
|
||||
input_parallel, self.weight, self.bias, name=self._name
|
||||
)
|
||||
|
||||
if self.gather_output and self.is_mp:
|
||||
output = mp_ops._c_concat(
|
||||
output_parallel, group=self.model_parallel_group
|
||||
)
|
||||
else:
|
||||
output = output_parallel
|
||||
return output
|
||||
|
||||
def sharded_state_dict(
|
||||
self,
|
||||
structured_name_prefix: str = "",
|
||||
):
|
||||
state_dict = self.state_dict(structured_name_prefix="")
|
||||
return build_sharded_state_dict(
|
||||
state_dict, {"weight": 1, "bias": 0}, structured_name_prefix
|
||||
)
|
||||
|
||||
|
||||
class MPScale(PyLayer):
|
||||
@staticmethod
|
||||
def forward(ctx, x, mp_degree):
|
||||
out = paddle.scale(x, 1.0 / mp_degree)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dout):
|
||||
return dout
|
||||
|
||||
|
||||
class RowParallelLinear(paddle.nn.Layer):
|
||||
"""Linear layer with mp parallelized(row).
|
||||
this class is used for splitting Linear Layer in mp group, row split the weight of the Linear layer.
|
||||
|
||||
Args:
|
||||
in_features(int): The number of input units.
|
||||
out_features(int): The number of output units.
|
||||
weight_attr(ParamAttr|None): The attribute for the learnable weight of this layer. The default value is None
|
||||
and the weight will be initialized to zero. For detailed information, please refer to paddle.ParamAttr.
|
||||
has_bias(bool): whether to add bias.
|
||||
input_is_parallel(bool): whether the input has already been split across the mp group.
|
||||
fuse_matmul_bias(bool): whether to fuse matmul and bias.
|
||||
mp_group(Group): The tensor parallel group.
|
||||
name(str, optional): Normally there is no need for user to set this parameter.
|
||||
For detailed information, please refer to :ref:`api_guide_Name` .
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed import fleet
|
||||
|
||||
>>> class SimpleMPNet(paddle.nn.Layer):
|
||||
... def __init__(self, vocab_size, hidden_size, inner_size, output_size):
|
||||
... super().__init__()
|
||||
... self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
|
||||
... hidden_size,
|
||||
... inner_size,
|
||||
... gather_output=False,
|
||||
... has_bias=True,
|
||||
... )
|
||||
... self.linear2 = fleet.meta_parallel.RowParallelLinear(
|
||||
... inner_size,
|
||||
... hidden_size,
|
||||
... input_is_parallel=True,
|
||||
... has_bias=True,
|
||||
... )
|
||||
... self.linear3 = paddle.nn.Linear(hidden_size, output_size)
|
||||
... self.embedding = fleet.meta_parallel.VocabParallelEmbedding(vocab_size, hidden_size)
|
||||
...
|
||||
... def forward(self, x):
|
||||
... x = self.embedding(x)
|
||||
... x = self.linear1(x)
|
||||
... x = self.linear2(x)
|
||||
... x = self.linear3(x)
|
||||
... return x
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features,
|
||||
out_features,
|
||||
weight_attr=None,
|
||||
has_bias=True,
|
||||
input_is_parallel=False,
|
||||
fuse_matmul_bias=False,
|
||||
mp_group=None,
|
||||
name=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.input_is_parallel = input_is_parallel
|
||||
self._weight_attr = weight_attr
|
||||
self._dtype = self._helper.get_default_dtype()
|
||||
self._name = name
|
||||
|
||||
self.model_parallel_group = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
|
||||
if mp_group is None
|
||||
else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
self.rank = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
|
||||
if mp_group is None
|
||||
else mp_group.rank
|
||||
)
|
||||
|
||||
self.is_mp = self.world_size > 1
|
||||
mp_configs = fleet.fleet._user_defined_strategy.hybrid_configs[
|
||||
"mp_configs"
|
||||
]
|
||||
self.mp_async_allreduce = self.is_mp and mp_configs.mp_async_allreduce
|
||||
self.mp_skip_c_identity = (
|
||||
self.is_mp
|
||||
and mp_configs.mp_async_allreduce
|
||||
and mp_configs.mp_skip_c_identity
|
||||
)
|
||||
self.mp_fused_linear_param_grad_add = (
|
||||
self.is_mp
|
||||
and mp_configs.mp_async_allreduce
|
||||
and mp_configs.mp_fused_linear_param_grad_add
|
||||
)
|
||||
if (
|
||||
self.mp_async_allreduce
|
||||
or self.mp_skip_c_identity
|
||||
or self.mp_fused_linear_param_grad_add
|
||||
):
|
||||
assert paddle.in_dynamic_mode(), (
|
||||
"mp_async_allreduce, mp_skip_c_identity and mp_fused_linear_param_grad_add are only available under dygraph mode"
|
||||
)
|
||||
assert in_features % self.world_size == 0, (
|
||||
f"Number of row of the weight for linear ({in_features}) must be"
|
||||
f" divisible by model parallel size ({self.world_size})"
|
||||
)
|
||||
|
||||
self.input_size_per_partition = in_features // self.world_size
|
||||
|
||||
if self.is_mp and paddle.in_dynamic_mode():
|
||||
with get_rng_state_tracker().rng_state():
|
||||
self.weight = self.create_parameter(
|
||||
shape=[self.input_size_per_partition, self.out_features],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
else:
|
||||
self.weight = self.create_parameter(
|
||||
shape=[self.input_size_per_partition, self.out_features],
|
||||
attr=self._weight_attr,
|
||||
dtype=self._dtype,
|
||||
is_bias=False,
|
||||
)
|
||||
|
||||
self.weight.is_distributed = True if self.is_mp else False
|
||||
if self.weight.is_distributed:
|
||||
self.weight.split_axis = 0
|
||||
|
||||
if has_bias:
|
||||
self.bias = self.create_parameter(
|
||||
shape=[self.out_features],
|
||||
attr=paddle.nn.initializer.Constant(value=0.0),
|
||||
dtype=self._dtype,
|
||||
is_bias=True,
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
self.linear = F.linear
|
||||
|
||||
if fuse_matmul_bias:
|
||||
if not is_fused_matmul_bias_supported():
|
||||
raise NotImplementedError(
|
||||
"You set fuse_matmul_bias=True in RowParallelLinear, "
|
||||
"however, the paddle you are using not support this operation. "
|
||||
"Please set fuse_matmul_bias=False or use paddle compiled "
|
||||
"with cuda 11.6 or higher."
|
||||
)
|
||||
from paddle.incubate.nn.functional import fused_linear
|
||||
|
||||
self.linear = fused_linear
|
||||
self.fuse_matmul_bias = fuse_matmul_bias
|
||||
|
||||
def forward(self, x):
|
||||
if self.input_is_parallel or (not self.is_mp):
|
||||
input_parallel = x
|
||||
else:
|
||||
# split last dim
|
||||
input_parallel = mp_ops._c_split(x, group=self.model_parallel_group)
|
||||
|
||||
if self.is_mp:
|
||||
if self.fuse_matmul_bias:
|
||||
bias = MPScale.apply(self.bias, self.world_size)
|
||||
output_parallel = self.linear(
|
||||
input_parallel, self.weight, bias, name=self._name
|
||||
)
|
||||
output = mp_ops._mp_allreduce(
|
||||
output_parallel,
|
||||
group=self.model_parallel_group,
|
||||
use_calc_stream=True,
|
||||
use_model_parallel=True,
|
||||
skip_c_identity_dynamic=self.mp_skip_c_identity,
|
||||
)
|
||||
else:
|
||||
output_parallel = self.linear(
|
||||
input_parallel, self.weight, name=self._name
|
||||
)
|
||||
output_ = mp_ops._mp_allreduce(
|
||||
output_parallel,
|
||||
group=self.model_parallel_group,
|
||||
use_calc_stream=True,
|
||||
use_model_parallel=True,
|
||||
skip_c_identity_dynamic=self.mp_skip_c_identity,
|
||||
)
|
||||
output = (
|
||||
output_ + self.bias if self.bias is not None else output_
|
||||
)
|
||||
else:
|
||||
output = self.linear(
|
||||
input_parallel, self.weight, self.bias, name=self._name
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
def sharded_state_dict(
|
||||
self,
|
||||
structured_name_prefix: str = "",
|
||||
):
|
||||
state_dict = self.state_dict(structured_name_prefix="")
|
||||
return build_sharded_state_dict(
|
||||
state_dict, {"weight": 0}, structured_name_prefix
|
||||
)
|
||||
|
||||
|
||||
class ParallelCrossEntropy(paddle.nn.Layer):
|
||||
"""CrossEntropy with mp parallelized.
|
||||
this class is used for splitting softmax cross entropy in mp group.
|
||||
|
||||
Args:
|
||||
mp_group(Group): The tensor parallel group.
|
||||
name(str, optional): Normally there is no need for user to set this parameter.
|
||||
For detailed information, please refer to :ref:`api_guide_Name` .
|
||||
ignore_index (long int, optional): Specifies a target value that is ignored and
|
||||
does not contribute to the loss. A negative value means that no label value
|
||||
needs to be ignored. Default is -100 .
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No img to demonstrate')
|
||||
>>> from paddle.distributed.fleet.layers.mpu import ParallelCrossEntropy
|
||||
>>> loss_func = ParallelCrossEntropy
|
||||
>>> loss = loss_func(img, label)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, mp_group=None, name=None, ignore_index=-100):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
self.model_parallel_group = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
|
||||
if mp_group is None
|
||||
else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
self.rank = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
|
||||
if mp_group is None
|
||||
else mp_group.rank
|
||||
)
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
def forward(self, input, label):
|
||||
loss = mp_ops._c_softmax_with_cross_entropy(
|
||||
input,
|
||||
label,
|
||||
group=self.model_parallel_group,
|
||||
ignore_index=self.ignore_index,
|
||||
)
|
||||
return loss
|
||||
|
||||
|
||||
class ParallelMultiLabelCrossEntropy(paddle.nn.Layer):
|
||||
"""CrossEntropy with mp parallelized.
|
||||
this class is used for splitting softmax cross entropy in mp group.
|
||||
|
||||
Args:
|
||||
mp_group(Group): The tensor parallel group.
|
||||
name(str, optional): Normally there is no need for user to set this parameter.
|
||||
For detailed information, please refer to :ref:`api_guide_Name` .
|
||||
ignore_index (long int, optional): Specifies a target value that is ignored and
|
||||
does not contribute to the loss. A negative value means that no label value
|
||||
needs to be ignored. Default is -100 .
|
||||
sum_multi_label_loss (bool, optional): Whether to sum the loss. Default is True .
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No img to demonstrate')
|
||||
>>> from paddle.distributed.fleet.layers.mpu import ParallelMultiLabelCrossEntropy
|
||||
>>> loss_func = ParallelMultiLabelCrossEntropy()
|
||||
>>> loss = loss_func(img, label, smooth_weight)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mp_group=None,
|
||||
name=None,
|
||||
ignore_index=-100,
|
||||
sum_multi_label_loss=True,
|
||||
):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
self.model_parallel_group = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group()
|
||||
if mp_group is None
|
||||
else mp_group
|
||||
)
|
||||
self.world_size = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size()
|
||||
if mp_group is None
|
||||
else mp_group.nranks
|
||||
)
|
||||
self.rank = (
|
||||
tp._HYBRID_PARALLEL_GROUP.get_model_parallel_rank()
|
||||
if mp_group is None
|
||||
else mp_group.rank
|
||||
)
|
||||
self.ignore_index = ignore_index
|
||||
self.sum_multi_label_loss = sum_multi_label_loss
|
||||
|
||||
def forward(self, input, label, smooth_weight):
|
||||
loss = mp_ops._c_softmax_with_multi_label_cross_entropy(
|
||||
input,
|
||||
label,
|
||||
smooth_weight,
|
||||
group=self.model_parallel_group,
|
||||
ignore_index=self.ignore_index,
|
||||
sum_multi_label_loss=self.sum_multi_label_loss,
|
||||
)
|
||||
return loss
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,266 @@
|
||||
# Copyright (c) 2022 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 contextlib
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import _legacy_C_ops
|
||||
from paddle.base import core
|
||||
from paddle.base.data_feeder import check_variable_and_dtype
|
||||
from paddle.common_ops_import import Variable
|
||||
from paddle.framework import LayerHelper, in_dynamic_mode
|
||||
|
||||
__all__ = []
|
||||
|
||||
MODEL_PARALLEL_RNG = 'model_parallel_rng'
|
||||
|
||||
# This file is inspired by Megatron to control random states for MP:
|
||||
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/random.py
|
||||
|
||||
|
||||
class RNGStatesTracker:
|
||||
"""
|
||||
Tracker the RNG states.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Map from name to the rng state.
|
||||
self.states_ = {}
|
||||
self.seeds_ = set()
|
||||
|
||||
def reset(self):
|
||||
self.states_ = {}
|
||||
self.seeds_ = set()
|
||||
|
||||
def add(self, name, seed):
|
||||
if seed in self.seeds_:
|
||||
raise ValueError(f'seed {seed} already exists')
|
||||
self.seeds_.add(seed)
|
||||
if name in self.states_:
|
||||
raise ValueError(f'state {name} already exists')
|
||||
orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
|
||||
# register a new state and set that state with the seed, store the indices into states_
|
||||
self.states_[name] = paddle.incubate.register_rng_state_as_index()
|
||||
paddle.seed(seed)
|
||||
paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
|
||||
|
||||
def get_states_tracker(self):
|
||||
states = {}
|
||||
orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
|
||||
for name in self.states_:
|
||||
# switch index to name
|
||||
paddle.incubate.set_rng_state(self.states_[name], use_index=True)
|
||||
# export the saved state
|
||||
states[name] = paddle.get_rng_state()
|
||||
paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
|
||||
return states
|
||||
|
||||
def set_states_tracker(self, states):
|
||||
orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
|
||||
for name in states:
|
||||
if name not in self.states_:
|
||||
raise ValueError(f'state {name} does not exists')
|
||||
# switch index to name
|
||||
paddle.incubate.set_rng_state(self.states_[name], use_index=True)
|
||||
# set the state to saved state
|
||||
paddle.set_rng_state(states[name])
|
||||
|
||||
paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def rng_state(self, name=MODEL_PARALLEL_RNG):
|
||||
if name not in self.states_:
|
||||
raise ValueError(f'state {name} does not exist')
|
||||
orig_rng_state_index = paddle.incubate.get_rng_state(use_index=True)
|
||||
paddle.incubate.set_rng_state(self.states_[name], use_index=True)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.states_[name] = paddle.incubate.get_rng_state(use_index=True)
|
||||
paddle.incubate.set_rng_state(orig_rng_state_index, use_index=True)
|
||||
|
||||
|
||||
RNG_STATE_TRACKER = RNGStatesTracker()
|
||||
|
||||
|
||||
def get_rng_state_tracker():
|
||||
return RNG_STATE_TRACKER
|
||||
|
||||
|
||||
def model_parallel_random_seed(seed=None):
|
||||
from paddle.distributed import fleet
|
||||
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
|
||||
mp_rank = hcg.get_model_parallel_rank()
|
||||
mp_size = hcg.get_model_parallel_world_size()
|
||||
|
||||
pp_rank = hcg.get_stage_id()
|
||||
pp_size = hcg.get_pipe_parallel_world_size()
|
||||
|
||||
if seed:
|
||||
global_seed = seed
|
||||
# dp/sharding seed is same
|
||||
local_seed = seed + 1 + mp_rank * pp_size + pp_rank
|
||||
else:
|
||||
global_seed = np.random.randint(0, 10000)
|
||||
local_seed = global_seed + 1 + mp_rank * pp_size + pp_rank
|
||||
|
||||
RNG_STATE_TRACKER.reset()
|
||||
RNG_STATE_TRACKER.add(MODEL_PARALLEL_RNG, local_seed)
|
||||
paddle.seed(global_seed)
|
||||
|
||||
|
||||
def dropout(
|
||||
x,
|
||||
p=0.5,
|
||||
axis=None,
|
||||
rng_name=None,
|
||||
training=True,
|
||||
mode="upscale_in_train",
|
||||
name=None,
|
||||
):
|
||||
"""
|
||||
Dropout is a regularization technique for reducing overfitting by preventing
|
||||
neuron co-adaption during training. The dropout operator randomly sets the
|
||||
outputs of some units to zero, while upscale others according to the given
|
||||
dropout probability.
|
||||
|
||||
Args:
|
||||
x (Tensor): The input tensor. The data type is float32 or float64.
|
||||
p (float|int): Probability of setting units to zero. Default 0.5.
|
||||
axis (int|list|tuple): The axis along which the dropout is performed. Default None.
|
||||
rng_name (str): The random seed generator name, which used to obtain deterministic results.
|
||||
training (bool): A flag indicating whether it is in train phrase or not. Default True.
|
||||
mode(str): ['upscale_in_train'(default) | 'downscale_in_infer'].
|
||||
|
||||
1. upscale_in_train(default), upscale the output at training time
|
||||
|
||||
- train: out = input * mask / ( 1.0 - dropout_prob )
|
||||
- inference: out = input
|
||||
|
||||
2. downscale_in_infer, downscale the output at inference
|
||||
|
||||
- train: out = input * mask
|
||||
- inference: out = input * (1.0 - dropout_prob)
|
||||
name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
|
||||
|
||||
Returns:
|
||||
A Tensor representing the dropout, has same shape and data type as `x` .
|
||||
|
||||
|
||||
Examples:
|
||||
We use ``p=0.5`` in the following description for simplicity.
|
||||
|
||||
1. When ``axis=None`` , this is commonly used dropout, which dropout each element of x randomly.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
Let's see a simple case when x is a 2d tensor with shape 2*3:
|
||||
[[1 2 3]
|
||||
[4 5 6]]
|
||||
we generate mask with the same shape as x, which is 2*3. The value of mask is
|
||||
sampled from a Bernoulli distribution randomly. For example, we may get such mask:
|
||||
[[0 1 0]
|
||||
[1 0 1]]
|
||||
So the output is obtained from elementwise multiply of x and mask:
|
||||
[[0 2 0]
|
||||
[4 0 6]]
|
||||
Using default setting, i.e. ``mode='upscale_in_train'`` ,
|
||||
if in training phase, the final upscale output is:
|
||||
[[0 4 0 ]
|
||||
[8 0 12]]
|
||||
if in test phase, the output is the same as input:
|
||||
[[1 2 3]
|
||||
[4 5 6]]
|
||||
we can also set ``mode='downscale_in_infer'`` , then
|
||||
if in training phase, the final output is:
|
||||
[[0 2 0]
|
||||
[4 0 6]]
|
||||
if in test phase, the scale output is:
|
||||
[[0.5 1. 1.5]
|
||||
[2. 2.5 3. ]]
|
||||
|
||||
"""
|
||||
if rng_name is None:
|
||||
return paddle.nn.functional.dropout(x, p, axis, training, mode, name)
|
||||
|
||||
if not isinstance(p, (float, int, Variable)):
|
||||
raise TypeError("p argument should be a number(int|float) or Variable")
|
||||
|
||||
# fast return for p == 0
|
||||
if isinstance(p, (int, float)) and p == 0:
|
||||
return x
|
||||
|
||||
assert 0 <= p <= 1, ValueError("p argument should between 0 and 1")
|
||||
assert mode in ('downscale_in_infer', 'upscale_in_train'), ValueError(
|
||||
"mode argument should be 'downscale_in_infer' or 'upscale_in_train'"
|
||||
)
|
||||
|
||||
assert axis is None, TypeError(
|
||||
"unsupported axis when using random seed generator"
|
||||
)
|
||||
|
||||
mode = (
|
||||
'downgrade_in_infer' if mode == 'downscale_in_infer' else mode
|
||||
) # semantic transfer
|
||||
|
||||
# dygraph using tracker, doesn't need determinate seed
|
||||
if in_dynamic_mode():
|
||||
out, mask = _legacy_C_ops.dropout(
|
||||
x,
|
||||
'dropout_prob',
|
||||
p,
|
||||
'is_test',
|
||||
not training,
|
||||
'fix_seed',
|
||||
False,
|
||||
'seed',
|
||||
0,
|
||||
'dropout_implementation',
|
||||
mode,
|
||||
)
|
||||
return out
|
||||
else:
|
||||
if isinstance(p, Variable) and not p.shape != [1]:
|
||||
raise TypeError(
|
||||
f"Required p.shape == [1] if type(p) is Variable, but received p.shape = {p.shape}"
|
||||
)
|
||||
|
||||
helper = LayerHelper('dropout', **locals())
|
||||
check_variable_and_dtype(
|
||||
x, 'x', ['float16', 'float32', 'float64'], 'dropout'
|
||||
)
|
||||
|
||||
seed = helper.create_variable_for_type_inference(dtype=paddle.int32)
|
||||
helper.append_op(type='seed', outputs={'Out': seed})
|
||||
|
||||
out = helper.create_variable_for_type_inference(dtype=x.dtype)
|
||||
mask = helper.create_variable_for_type_inference(
|
||||
dtype=core.VarDesc.VarType.UINT8, stop_gradient=True
|
||||
)
|
||||
|
||||
helper.append_op(
|
||||
type='dropout',
|
||||
inputs={'X': [x], 'Seed': seed},
|
||||
outputs={'Out': [out], 'Mask': [mask]},
|
||||
attrs={
|
||||
'dropout_prob': p,
|
||||
'is_test': not training,
|
||||
'dropout_implementation': mode,
|
||||
},
|
||||
)
|
||||
return out
|
||||
Reference in New Issue
Block a user