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/nn/functional/__init__.py
+329
View File
@@ -0,0 +1,329 @@
# Copyright (c) 2020 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.
# TODO: import all neural network related api under this directory,
# including layers, linear, conv, rnn etc.
from .activation import (
celu,
elu,
elu_,
gelu,
glu,
gumbel_softmax,
hardshrink,
hardsigmoid,
hardswish,
hardtanh,
hardtanh_,
leaky_relu,
leaky_relu_,
log_sigmoid,
log_softmax,
maxout,
mish,
prelu,
relu,
relu6,
relu_,
rrelu,
selu,
sigmoid,
silu,
softmax,
softmax_,
softplus,
softshrink,
softsign,
swiglu,
swish,
tanh,
tanh_,
tanhshrink,
thresholded_relu,
thresholded_relu_,
)
from .common import (
alpha_dropout,
bilinear,
class_center_sample,
cosine_similarity,
dropout,
dropout1d,
dropout2d,
dropout3d,
feature_alpha_dropout,
fold,
interpolate,
label_smooth,
linear,
pad,
unfold,
upsample,
zeropad2d,
)
from .conv import (
conv1d,
conv1d_transpose,
conv2d,
conv2d_transpose,
conv3d,
conv3d_transpose,
)
from .distance import pairwise_distance, pdist # noqa: F401
from .extension import (
diag_embed, # noqa: F401
gather_tree,
sequence_mask,
temporal_shift,
)
from .flash_attention import (
flash_attention_v3_varlen,
flash_attn_qkvpacked,
flash_attn_varlen_qkvpacked,
flashmask_attention,
flashmask_get_unique_id,
sdp_kernel, # noqa: F401
)
from .input import (
embedding,
embedding_renorm_, # noqa: F401
one_hot,
)
from .loss import (
adaptive_log_softmax_with_loss,
binary_cross_entropy,
binary_cross_entropy_with_logits,
cosine_embedding_loss,
cross_entropy,
ctc_loss,
dice_loss,
gaussian_nll_loss,
hinge_embedding_loss,
hsigmoid_loss,
kl_div,
l1_loss,
log_loss,
margin_cross_entropy,
margin_ranking_loss,
mse_loss,
multi_label_margin_loss,
multi_label_soft_margin_loss,
multi_margin_loss,
nll_loss,
npair_loss,
poisson_nll_loss,
rnnt_loss,
sigmoid_focal_loss,
smooth_l1_loss,
soft_margin_loss,
softmax_with_cross_entropy,
square_error_cost,
triplet_margin_loss,
triplet_margin_with_distance_loss,
)
from .moe_permute import moe_permute
from .moe_unpermute import moe_unpermute
from .norm import (
batch_norm,
group_norm,
instance_norm,
layer_norm,
local_response_norm,
normalize,
rms_norm,
)
from .pooling import (
adaptive_avg_pool1d,
adaptive_avg_pool2d,
adaptive_avg_pool3d,
adaptive_max_pool1d,
adaptive_max_pool2d,
adaptive_max_pool3d,
avg_pool1d,
avg_pool2d,
avg_pool3d,
fractional_max_pool2d,
fractional_max_pool3d,
lp_pool1d,
lp_pool2d,
max_pool1d,
max_pool2d,
max_pool3d,
max_unpool1d,
max_unpool2d,
max_unpool3d,
)
from .sdpa import scaled_dot_product_attention
from .sparse_attention import sparse_attention
from .vision import (
affine_grid,
channel_shuffle,
grid_sample,
pixel_shuffle,
pixel_unshuffle,
)
logsigmoid = log_sigmoid
conv_transpose1d = conv1d_transpose
conv_transpose2d = conv2d_transpose
conv_transpose3d = conv3d_transpose
huber_loss = smooth_l1_loss
multilabel_margin_loss = multi_label_margin_loss
multilabel_soft_margin_loss = multi_label_soft_margin_loss
__all__ = [
'celu',
'conv1d',
'conv1d_transpose',
'conv2d',
'conv2d_transpose',
'conv3d',
'conv3d_transpose',
'conv_transpose1d',
'conv_transpose2d',
'conv_transpose3d',
'pairwise_distance',
'elu',
'elu_',
'gelu',
'hardshrink',
'hardtanh',
'hardtanh_',
'hardsigmoid',
'hardswish',
'leaky_relu',
'leaky_relu_',
'log_sigmoid',
'logsigmoid',
'maxout',
'prelu',
'relu',
'relu_',
'relu6',
'selu',
'softmax',
'softmax_',
'softplus',
'softshrink',
'softsign',
'sigmoid',
'silu',
'swiglu',
'swish',
'mish',
'tanh',
'tanh_',
'tanhshrink',
'thresholded_relu',
'thresholded_relu_',
'log_softmax',
'glu',
'gumbel_softmax',
'sequence_mask',
'dropout',
'dropout1d',
'dropout2d',
'dropout3d',
'alpha_dropout',
'feature_alpha_dropout',
'label_smooth',
'linear',
'pad',
'zeropad2d',
'unfold',
'interpolate',
'upsample',
'bilinear',
'cosine_similarity',
'avg_pool1d',
'avg_pool2d',
'avg_pool3d',
'lp_pool1d',
'lp_pool2d',
'max_pool1d',
'max_pool2d',
'max_pool3d',
'max_unpool1d',
'max_unpool2d',
'max_unpool3d',
'moe_permute',
'moe_unpermute',
'adaptive_avg_pool1d',
'adaptive_avg_pool2d',
'adaptive_avg_pool3d',
'adaptive_max_pool1d',
'adaptive_max_pool2d',
'adaptive_max_pool3d',
'fractional_max_pool2d',
'fractional_max_pool3d',
'binary_cross_entropy',
'binary_cross_entropy_with_logits',
'cross_entropy',
'dice_loss',
'hsigmoid_loss',
'kl_div',
'l1_loss',
'log_loss',
'mse_loss',
'margin_ranking_loss',
'multi_label_soft_margin_loss',
'nll_loss',
'poisson_nll_loss',
'npair_loss',
'sigmoid_focal_loss',
'smooth_l1_loss',
'softmax_with_cross_entropy',
'margin_cross_entropy',
'square_error_cost',
'ctc_loss',
'rnnt_loss',
'hinge_embedding_loss',
'affine_grid',
'grid_sample',
'local_response_norm',
'pixel_shuffle',
'pixel_unshuffle',
'channel_shuffle',
'embedding',
'gather_tree',
'one_hot',
'normalize',
'temporal_shift',
'batch_norm',
'layer_norm',
'rms_norm',
'instance_norm',
'class_center_sample',
'sparse_attention',
'fold',
'cosine_embedding_loss',
'rrelu',
'triplet_margin_with_distance_loss',
'triplet_margin_loss',
'adaptive_log_softmax_with_loss',
'multi_margin_loss',
'multi_label_margin_loss',
'multilabel_margin_loss',
'multilabel_soft_margin_loss',
'soft_margin_loss',
'gaussian_nll_loss',
'scaled_dot_product_attention',
'flashmask_attention',
'flashmask_get_unique_id',
'flash_attn_qkvpacked',
"flash_attention_v3_varlen",
'flash_attn_varlen_qkvpacked',
'group_norm',
'huber_loss',
]
python/paddle/nn/functional/activation.py
+2048
View File
File diff suppressed because it is too large Load Diff
python/paddle/nn/functional/common.py
+3053
View File
File diff suppressed because it is too large Load Diff
python/paddle/nn/functional/conv.py
+1981
View File
File diff suppressed because it is too large Load Diff
python/paddle/nn/functional/distance.py
+163
View File
@@ -0,0 +1,163 @@
# 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 __future__ import annotations
import paddle
from paddle import _C_ops
from paddle.framework import in_dynamic_or_pir_mode
from paddle.utils.decorator_utils import ParamAliasDecorator
from ...base.data_feeder import check_type, check_variable_and_dtype
from ...base.layer_helper import LayerHelper
__all__ = []
@ParamAliasDecorator(
{
"x": ["x1"],
"y": ["x2"],
"epsilon": ["eps"],
}
)
def pairwise_distance(
x: paddle.Tensor,
y: paddle.Tensor,
p: float = 2.0,
epsilon: float = 1e-6,
keepdim: bool = False,
name: str | None = None,
) -> paddle.Tensor:
r"""
It computes the pairwise distance between two vectors. The
distance is calculated by p-order norm:
.. math::
\Vert x \Vert _p = \left( \sum_{i=1}^n \vert x_i \vert ^ p \right) ^ {1/p}.
Parameters:
x (Tensor): Tensor, shape is :math:`[N, D]` or :math:`[D]`, where :math:`N`
is batch size, :math:`D` is the dimension of vector. Available dtype is
float16, float32, float64.
y (Tensor): Tensor, shape is :math:`[N, D]` or :math:`[D]`, where :math:`N`
is batch size, :math:`D` is the dimension of vector. Available dtype is
float16, float32, float64.
p (float, optional): The order of norm. Default: :math:`2.0`.
epsilon (float, optional): Add small value to avoid division by zero.
Default: :math:`1e-6`.
keepdim (bool, optional): Whether to reserve the reduced dimension
in the output Tensor. The result tensor is one dimension less than
the result of ``|x-y|`` unless :attr:`keepdim` is True. Default: False.
name (str|None, optional): For details, please refer to :ref:`api_guide_Name`.
Generally, no setting is required. Default: None.
Returns:
Tensor, the dtype is same as input tensor.
- If :attr:`keepdim` is True, the output shape is :math:`[N, 1]` or :math:`[1]`,
depending on whether the input has data shaped as :math:`[N, D]`.
- If :attr:`keepdim` is False, the output shape is :math:`[N]` or :math:`[]`,
depending on whether the input has data shaped as :math:`[N, D]`.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.to_tensor([[1.0, 3.0], [3.0, 5.0]], dtype=paddle.float64)
>>> y = paddle.to_tensor([[5.0, 6.0], [7.0, 8.0]], dtype=paddle.float64)
>>> distance = paddle.nn.functional.pairwise_distance(x, y)
>>> print(distance)
Tensor(shape=[2], dtype=float64, place=Place(cpu), stop_gradient=True,
[4.99999860, 4.99999860])
"""
if in_dynamic_or_pir_mode():
sub = _C_ops.subtract(x, y)
# p_norm op has not used epsilon, so change it to the following.
if epsilon != 0.0:
epsilon = paddle.to_tensor([epsilon], dtype=sub.dtype)
sub = _C_ops.add(sub, epsilon)
return _C_ops.p_norm(sub, p, -1, 0.0, keepdim, False)
else:
check_type(p, 'porder', (float, int), 'PairwiseDistance')
check_type(epsilon, 'epsilon', (float), 'PairwiseDistance')
check_type(keepdim, 'keepdim', (bool), 'PairwiseDistance')
check_variable_and_dtype(
x, 'x', ['float16', 'float32', 'float64'], 'PairwiseDistance'
)
check_variable_and_dtype(
y, 'y', ['float16', 'float32', 'float64'], 'PairwiseDistance'
)
sub = paddle.subtract(x, y)
if epsilon != 0.0:
epsilon_var = sub.block.create_var(dtype=sub.dtype)
epsilon_var = paddle.full(
shape=[1], fill_value=epsilon, dtype=sub.dtype
)
sub = paddle.add(sub, epsilon_var)
helper = LayerHelper("PairwiseDistance", name=name)
attrs = {
'axis': -1,
'porder': p,
'keepdim': keepdim,
'epsilon': 0.0,
}
out = helper.create_variable_for_type_inference(dtype=x.dtype)
helper.append_op(
type='p_norm', inputs={'X': sub}, outputs={'Out': out}, attrs=attrs
)
return out
def pdist(
x: paddle.Tensor, p: float = 2.0, name: str | None = None
) -> paddle.Tensor:
r'''
Computes the p-norm distance between every pair of row vectors in the input.
Args:
x (Tensor): The input tensor with shape :math:`N \times M`.
p (float, optional): The value for the p-norm distance to calculate between each vector pair. Default: :math:`2.0`.
name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None.
Returns:
Tensor with shape :math:`N(N-1)/2` , the dtype is same as input tensor.
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.seed(2023)
>>> a = paddle.randn([4, 5])
>>> print(a)
Tensor(shape=[4, 5], dtype=float32, place=Place(cpu), stop_gradient=True,
[[ 0.06132207, 1.11349595, 0.41906244, -0.24858207, -1.85169315],
[-1.50370061, 1.73954511, 0.13331604, 1.66359663, -0.55764782],
[-0.59911072, -0.57773495, -1.03176904, -0.33741450, -0.29695082],
[-1.50258386, 0.67233968, -1.07747352, 0.80170447, -0.06695852]])
>>> pdist_out = paddle.pdist(a)
>>> print(pdist_out)
Tensor(shape=[6], dtype=float32, place=Place(cpu), stop_gradient=True,
[2.87295413, 2.79758120, 3.02793980, 3.40844536, 1.89435327, 1.93171620])
'''
x_shape = list(x.shape)
assert len(x_shape) == 2, "The x must be 2-dimensional"
d = paddle.linalg.norm(x[..., None, :] - x[..., None, :, :], p=p, axis=-1)
mask = ~paddle.tril(paddle.ones(d.shape, dtype='bool'))
return paddle.masked_select(d, mask)
python/paddle/nn/functional/extension.py
+348
View File
@@ -0,0 +1,348 @@
# Copyright (c) 2020 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
from typing import (
TYPE_CHECKING,
)
from paddle import _C_ops, tensor
from paddle.utils import deprecated
from ...base.data_feeder import (
check_dtype,
check_type,
check_variable_and_dtype,
)
from ...base.layer_helper import LayerHelper
from ...common_ops_import import Variable
from ...framework import (
convert_nptype_to_datatype_or_vartype,
core,
in_dynamic_or_pir_mode,
)
if TYPE_CHECKING:
from paddle import Tensor
from paddle._typing import DataLayout2D, DTypeLike
__all__ = []
@deprecated(
since="2.5.2",
update_to="paddle.diag_embed",
level=1,
reason="diag_embed in paddle.nn.functional will be removed in future",
)
def diag_embed(
input: Tensor, offset: int = 0, dim1: int = -2, dim2: int = -1
) -> Tensor:
return tensor.diag_embed(input, offset, dim1, dim2)
def sequence_mask(
x: Tensor,
maxlen: int | None = None,
dtype: DTypeLike = 'int64',
name: str | None = None,
) -> Tensor:
r"""
**SequenceMask Layer**
This layer outputs a mask according to the input :code:`x` and
:code:`maxlen` with data type of :code:`dtype`.
Supposing :code:`x` is a Tensor with shape [d_1, d_2, ..., d_n], the
:code:`y` is a mask with shape [d_1, d_2, ..., d_n, maxlen], where:
.. math::
y(i_1, i_2,..., i_n, j) = (j < x(i_1, i_2,..., i_n))
.. code-block:: text
Case:
Consider input:
x = [3, 1, 1, 0] max_len = 4
then we get out:
mask = [[1, 1, 1, 0],
[1, 0, 0, 0],
[1, 0, 0, 0],
[0, 0, 0, 0]]
Args:
x (Variable): Input tensor of sequence_mask layer, \
whose elements are integers less than :code:`maxlen`. \
Tensor with shape [d_1, d_2, ..., d_n].
maxlen (int|None, optional): Maximum length of the sequence. If :code:`maxlen` \
is None, it would be replace with :math:`max(x)`.
dtype (np.dtype|paddle.dtype|str, optional): Data type of the output, \
``int64`` by default.
name(str|None, optional): For detailed information, please refer \
to :ref:`api_guide_Name`. Usually name is no need to set and \
None by default.
Returns:
Tensor, The output sequence mask. Tensor with shape [d_1, d_2, ..., d_n, maxlen] \
and data type of :code:`dtype`. The data type should be bool, float32, float64, int8, \
int32 or int64.
Examples:
.. code-block:: pycon
>>> import paddle
>>> lengths = paddle.to_tensor([10, 9, 8])
>>> mask = paddle.nn.functional.sequence_mask(lengths)
>>> print(mask)
Tensor(shape=[3, 10], dtype=int64, place=Place(cpu), stop_gradient=True,
[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
"""
if in_dynamic_or_pir_mode():
if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
dtype = convert_nptype_to_datatype_or_vartype(dtype)
if maxlen is None:
maxlen = -1
out = _C_ops.sequence_mask(x, maxlen, dtype)
out.stop_gradient = True
return out
helper = LayerHelper('sequence_mask', **locals())
out = helper.create_variable_for_type_inference(dtype=dtype)
inputs = {'X': [x]}
attrs = {'out_dtype': out.dtype}
if maxlen is not None:
if isinstance(maxlen, Variable):
inputs['MaxLenTensor'] = maxlen
else:
attrs['maxlen'] = maxlen
helper.append_op(
type='sequence_mask', inputs=inputs, outputs={'Y': out}, attrs=attrs
)
out.stop_gradient = True
return out
def gather_tree(ids: Tensor, parents: Tensor) -> Tensor:
r"""
To be used after beam search. After beam search, we get selected ids at
each time step and the corresponding parents in the search tree. Both ids
and parents have the layout :attr:`[max_time, batch_size, beam_size]`. Then
:attr:`gather_tree` is used to backtrace from the last time step and
generate the full sequences by collecting selected ids.
Here is an example:
.. code-block:: text
Given:
ids = [[[2 2]
[6 1]]
[[3 9]
[6 1]]
[[0 1]
[9 0]]]
parents = [[[0 0]
[1 1]]
[[1 0]
[1 0]]
[[0 0]
[0 1]]]
Then:
gather_tree(ids, parents)
= [[[2 2]
[1 6]]
[[3 3]
[6 1]]
[[0 1]
[9 0]]]
Args:
ids(Tensor): A Tensor with shape :attr:`[length, batch_size, beam_size]`
and data type :attr:`int32` or :attr:`int64`. It contains the selected
ids of all time steps.
parents(Tensor): A Tensor with the same shape and data type as :attr:`ids`,
It contains the parents corresponding to selected ids when searching
among beams.
Returns:
A Tensor with the same shape and data type as :attr:`ids`. \
It contains the full sequences. The sequences are collected from \
:attr:`ids` by backtracing according to :attr:`parents`.
Examples:
.. code-block:: pycon
>>> import paddle
>>> ids = paddle.to_tensor([[[2, 2], [6, 1]], [[3, 9], [6, 1]], [[0, 1], [9, 0]]])
>>> parents = paddle.to_tensor([[[0, 0], [1, 1]], [[1, 0], [1, 0]], [[0, 0], [0, 1]]])
>>> final_sequences = paddle.nn.functional.gather_tree(ids, parents)
>>> [[[2, 2], [1, 6]], [[3, 3], [6, 1]], [[0, 1], [9, 0]]]
>>> final_sequences = paddle.nn.functional.gather_tree(ids, parents)
>>> print(final_sequences)
Tensor(shape=[3, 2, 2], dtype=int64, place=Place(cpu), stop_gradient=True,
[[[2, 2],
[1, 6]],
[[3, 3],
[6, 1]],
[[0, 1],
[9, 0]]])
"""
if ids.ndim != 3:
raise ValueError(
"The input ids must be a 3D tensor with shape [length, batch_size, beam_size]"
)
if ids.ndim != parents.ndim:
raise ValueError("The ids's shape must be the same as parents' shape. ")
if in_dynamic_or_pir_mode():
check_dtype(parents.dtype, "parents", ['int32', 'int64'], 'gather_tree')
return _C_ops.gather_tree(ids, parents)
else:
helper = LayerHelper('gather_tree', **locals())
check_variable_and_dtype(ids, 'ids', ['int32', 'int64'], 'gather_tree')
check_variable_and_dtype(
parents, 'parents', ['int32', 'int64'], 'gather_tree'
)
out = helper.create_variable_for_type_inference(dtype=ids.dtype)
helper.append_op(
type="gather_tree",
inputs={"Ids": ids, "Parents": parents},
outputs={"Out": out},
)
return out
def temporal_shift(
x: Tensor,
seg_num: int,
shift_ratio: float = 0.25,
name: str | None = None,
data_format: DataLayout2D | str = 'NCHW',
) -> Tensor:
"""
**Temporal Shift Operator**
Calculate the temporal shifting features for Input(X).
Input(X) should be in shape of [N*T, C, H, W] or [N*T, H, W, C], while
N is the batch size, T is the temporal segment number specified by
:attr:`seg_num`, C is the channel number, H and W is the height and
width of features.
Temporal Shifting is calculated as follows when data format is NCHW:
Step 1: Reshape Input(X) to [N, T, C, H, W].
Step 2: Pad 0 to reshaping result in the 2nd(T) dimension with
padding width as 1 on each side, padding result will be in shape
of [N, T+2, C, H, W].
Step 3: Assume :attr:`shift_ratio` is :math:`1/4`, slice padding
result as follows:
$$
slice1 = x[:, :T, :C/4, :, :]
$$
$$
slice2 = x[:, 2:T+2, C/4:C/2, :, :]
$$
$$
slice3 = x[:, 1:T+1, C/2:, :, :]
$$
Step 4: Concatenate three slices along the 3rd(C) dimension and
reshape result to [N*T, C, H, W].
For details of temporal shifting, please refer to paper:
`Temporal Shift Module <http://arxiv.org/abs/1811.08383>`_ .
Args:
x(Tensor): ${x_comment}
seg_num(int): ${seg_num_comment}
shift_ratio(float): ${shift_ratio_comment}
name(str|None, optional): For detailed information, please refer
to :ref:`api_guide_Name`. Usually name is no need to set and
None by default.
data_format(str, optional): Data format that specifies the layout of input.
It can be "NCHW" or "NHWC". Default: "NCHW".
Returns:
out(Tensor): The temporal shifting result is a tensor with the
same shape and same data type as the input.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn.functional as F
>>> input = paddle.randn([6, 4, 2, 2])
>>> out = F.temporal_shift(x=input, seg_num=2, shift_ratio=0.2)
"""
if data_format not in ["NCHW", "NHWC"]:
raise ValueError(
"Attr(data_format) should be 'NCHW' or 'NHWC'. "
f"Received Attr(data_format): {data_format}."
)
if in_dynamic_or_pir_mode():
return _C_ops.temporal_shift(x, seg_num, shift_ratio, data_format)
else:
helper = LayerHelper("temporal_shift", **locals())
check_variable_and_dtype(
x,
'x',
['float16', 'uint16', 'float32', 'float64'],
'temporal_shift',
)
check_type(seg_num, 'seg_num', int, 'temporal_shift')
check_type(shift_ratio, 'shift_ratio', float, 'temporal_shift')
out = helper.create_variable_for_type_inference(dtype=x.dtype)
if not isinstance(seg_num, int):
raise TypeError("seg_num must be int type.")
helper.append_op(
type="temporal_shift",
inputs={"X": x},
outputs={"Out": out},
attrs={
"seg_num": seg_num,
"shift_ratio": shift_ratio,
"data_format": data_format,
},
)
return out
python/paddle/nn/functional/flash_attention.py
+2304
View File
File diff suppressed because it is too large Load Diff
python/paddle/nn/functional/input.py
+363
View File
@@ -0,0 +1,363 @@
# Copyright (c) 2020 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
from typing import TYPE_CHECKING
import paddle
from paddle import _C_ops
from paddle.utils.decorator_utils import param_one_alias
from ...base.data_feeder import check_variable_and_dtype
from ...base.layer_helper import LayerHelper
from ...common_ops_import import Variable
from ...framework import in_dynamic_or_pir_mode
if TYPE_CHECKING:
from paddle import Tensor
__all__ = []
@param_one_alias(["x", "input"])
def one_hot(
x: Tensor,
num_classes: int = -1,
name: str | None = None,
) -> Tensor:
"""
The operator converts each id in the input `x` to an one-hot vector with a
`num_classes` length. The value in the vector dimension corresponding to the id
is 1, and the value in the remaining dimension is 0.
The shape of output Tensor is generated by appending `num_classes` dimension
behind the last dimension of the `x` shape.
.. code-block:: text
Example 1:
input:
x.shape = [4]
x.data = [1, 1, 3, 0]
num_classes = 4
output:
Out.shape = [4, 4]
Out.data = [[0., 1., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 0., 1.],
[1., 0., 0., 0.]]
Example 2:
input:
x.shape = [4]
x.data = [1, 1, 5, 0]
num_classes = 4
output: Throw an exception for Illegal value
The second dimension in X is 5, which is greater than num_classes,
so it throws an exception.
.. note::
Alias Support: The parameter name ``input`` can be used as an alias for ``x``.
For example, ``one_hot(input=tensor_x, ...)`` is equivalent to ``one_hot(x=tensor_x, ...)``.
Args:
x (Tensor): Tensor with shape :math:`[N_1, N_2, ..., N_k]` ,
which contains at least one dimension. The data type is int32 or int64.
Alias: ``input``.
num_classes (int): An integer defining the `num_classes` of the one hot dimension. If input `x`
is word id, `num_classes` is generally the dictionary size. Default value: -1.
name (str|None, optional): For detailed information, please refer
to :ref:`api_guide_Name`. Usually name is no need to set and
None by default.
Returns:
Tensor, The one-hot representations of `x`. A Tensor with type float32.
Examples:
.. code-block:: pycon
>>> import paddle
>>> # Correspond to the first example above, where label.shape is 4 and one_hot_label.shape is [4, 4].
>>> label = paddle.to_tensor([1, 1, 3, 0], dtype='int64')
>>> print(label.shape)
paddle.Size([4])
>>> one_hot_label = paddle.nn.functional.one_hot(label, num_classes=4)
>>> print(one_hot_label.shape)
paddle.Size([4, 4])
>>> print(one_hot_label)
Tensor(shape=[4, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
[[0., 1., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 0., 1.],
[1., 0., 0., 0.]])
"""
if not isinstance(num_classes, paddle.pir.Value) and num_classes == -1:
num_classes = x.max() + 1
if in_dynamic_or_pir_mode():
return _C_ops.one_hot(x, num_classes)
else:
check_variable_and_dtype(x, 'input', ['int32', 'int64'], 'one_hot_v2')
helper = LayerHelper("one_hot_v2", **locals())
one_hot_out = helper.create_variable_for_type_inference(dtype='float32')
if not isinstance(num_classes, Variable):
# user attribute
inputs = {'X': x}
attrs = {'depth': num_classes, 'allow_out_of_range': False}
else:
num_classes.stop_gradient = True
inputs = {'X': x, 'depth_tensor': num_classes}
attrs = {'allow_out_of_range': False}
helper.append_op(
type="one_hot_v2",
inputs=inputs,
attrs=attrs,
outputs={'Out': one_hot_out},
stop_gradient=True,
)
return one_hot_out
def embedding_renorm_(
x: Tensor, weight: Tensor, max_norm: float, norm_type: float = 2.0
) -> Tensor:
r"""
Renorm the weight of embedding with respect to the provided :attr:`max_norm` and :attr:`norm_type` .
Note:
In the dynamic graph mode, the input weight will be updated in-place, and the return value will be the changed weight.
Args:
x (Tensor): A Tensor with type int32/int64, which contains the id information. The value of the input id should
satisfy :math:`0<= id < weight.shape[0]` .
weight (Tensor): The weight. A Tensor with shape of lookup table parameter. It should have two elements which
indicates the size of the dictionary of embeddings and the size of each embedding vector respectively.
max_norm (float): The maximum norm for each embedding vector.
norm_type (float, optional): The p of the p-norm to compute for the max_norm option. Default: 2.0.
Returns:
Tensor, The updated weight. The data type is the same as :attr:`weight`.
"""
with paddle.set_grad_enabled(False):
unique_x = paddle.unique(x)
selected_rows = paddle.index_select(weight, unique_x)
norm = paddle.norm(selected_rows, p=norm_type, axis=1, keepdim=True)
mask = norm > max_norm
scale = max_norm / (norm + 1e-7)
scale = paddle.where(mask, scale, paddle.ones_like(scale))
scale = paddle.expand_as(scale, selected_rows)
updated_rows = selected_rows * scale
paddle.scatter_(weight, unique_x, updated_rows, overwrite=True)
return weight
@param_one_alias(["x", "input"])
def embedding(
x: Tensor,
weight: Tensor,
padding_idx: int | None = None,
max_norm: float | None = None,
norm_type: float = 2.0,
sparse: bool = False,
scale_grad_by_freq: bool = False,
name: str | None = None,
) -> Tensor:
r"""
Used to lookup embeddings vector of ids provided by :attr:`x` .
The shape of output Tensor is generated by appending the last dimension of the input Tensor shape
with embedding size.
Note:
The id in :attr:`x` must satisfy :math:`0 <= id < weight.shape[0]` ,
otherwise the program will throw an exception and exit.
.. code-block:: text
x is a Tensor.
padding_idx = -1
x.data = [[1, 3], [2, 4], [4, 127]]
x.shape = [3, 2]
weight.shape = [128, 16]
output is a Tensor:
out.shape = [3, 2, 16]
out.data = [[[0.129435295, 0.244512452, ..., 0.436322452],
[0.345421456, 0.524563927, ..., 0.144534654]],
[[0.345249859, 0.124939536, ..., 0.194353745],
[0.945345345, 0.435394634, ..., 0.435345365]],
[[0.945345345, 0.435394634, ..., 0.435345365],
[0.0, 0.0, ..., 0.0 ]]] # padding data
The input padding_idx is less than 0, it is automatically converted to padding_idx = -1 + 128 = 127
It will pad all-zero data when id is 127.
.. note::
Alias Support: The parameter name ``input`` can be used as an alias for ``x``.
For example, ``embedding(input=tensor_x, ...)`` is equivalent to ``embedding(x=tensor_x, ...)``.
Args:
x (Tensor): A Tensor with type int32/int64, which contains the id information. The value of the input id should
satisfy :math:`0 <= id < weight.shape[0]` .
Alias: ``input``.
weight (Tensor): The weight. A Tensor with shape of lookup table parameter. It should have two elements which
indicates the size of the dictionary of embeddings and the size of each embedding vector respectively.
sparse (bool, optional): The flag indicating whether to use sparse update. This parameter only
affects the performance of the backwards gradient update. It is recommended to set
True because sparse update is faster. But some optimizers does not support sparse update,
such as :ref:`api_paddle_optimizer_adadelta_Adadelta` , :ref:`api_paddle_optimizer_adamax_Adamax` , :ref:`api_paddle_optimizer_lamb_Lamb`.
In these cases, sparse must be False. Default: False.
padding_idx (int|None, optional): padding_idx needs to be in the interval [-weight.shape[0], weight.shape[0]).
If :math:`padding\_idx < 0`, the :math:`padding\_idx` will automatically be converted
to :math:`weight.shape[0] + padding\_idx` . It will output all-zero padding data whenever lookup
encounters :math:`padding\_idx` in id. And the padding data will not be updated while training.
If set None, it makes no effect to output. Default: None.
max_norm (float, optional): If provided, will renormalize the embedding vectors to have a norm larger than
:attr:`max\_norm` . It will inplace update the input embedding weight in dynamic graph mode. Default: None.
norm_type (float, optional): The p of the p-norm to compute for the max_norm option. Default: 2.0.
scale_grad_by_freq (bool, optional): Indicating whether to scale the gradients by the inverse frequency of the
word ids in input `x`. Default: False.
name (str|None, optional): For detailed information, please refer
to :ref:`api_guide_Name`. Usually name is no need to set and
None by default.
Returns:
Tensor, Embedding Tensor mapped by x. The data type is the same as :attr:`weight`.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn as nn
>>> x0 = paddle.arange(3, 6).reshape((3, 1)).astype(paddle.int64)
>>> w0 = paddle.full(shape=(10, 3), fill_value=2).astype(paddle.float32)
>>> x = paddle.to_tensor(x0, stop_gradient=False)
>>> print(x.numpy())
[[3]
[4]
[5]]
>>> print(x.shape)
paddle.Size([3, 1])
>>> w = paddle.to_tensor(w0, stop_gradient=False)
>>> print(w.numpy())
[[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]]
>>> print(w.shape)
paddle.Size([10, 3])
>>> emb = nn.functional.embedding(x=x, weight=w, sparse=True, name="embedding")
>>> print(emb.numpy())
[[[2. 2. 2.]]
[[2. 2. 2.]]
[[2. 2. 2.]]]
>>> print(emb.shape)
paddle.Size([3, 1, 3])
"""
padding_idx = (
-1
if padding_idx is None
else (
padding_idx if padding_idx >= 0 else (weight.shape[0] + padding_idx)
)
)
if weight.shape[0] != 0 and (
padding_idx >= weight.shape[0] or padding_idx < -weight.shape[0]
):
raise ValueError(
f"padding_idx must be within [-{weight.shape[0]}, {weight.shape[0]})"
)
if max_norm and weight.size != 0:
weight = embedding_renorm_(
x, weight, max_norm=max_norm, norm_type=norm_type
)
if scale_grad_by_freq:
if sparse:
raise AttributeError(
"scale_grad_by_freq = True is not supported with sparse update."
)
if in_dynamic_or_pir_mode():
return _C_ops.embedding_with_scaled_gradient(x, weight, padding_idx)
else:
helper = LayerHelper('embedding_with_scaled_gradient', **locals())
dtype = helper.input_dtype(input_param_name='weight')
check_variable_and_dtype(
x,
'input',
['uint8', 'int8', 'int16', 'int32', 'int64'],
'embedding_with_scaled_gradient',
)
tmp = helper.create_variable_for_type_inference(dtype)
helper.append_op(
type='embedding_with_scaled_gradient',
inputs={'x': x, 'weight': weight},
outputs={'out': tmp},
attrs={'padding_idx': padding_idx},
)
return tmp
else:
if in_dynamic_or_pir_mode():
return _C_ops.embedding(x, weight, padding_idx, sparse)
else:
helper = LayerHelper('embedding', **locals())
dtype = helper.input_dtype(input_param_name='weight')
check_variable_and_dtype(
x,
'input',
['uint8', 'int8', 'int16', 'int32', 'int64'],
'embedding',
)
is_distributed = False
remote_prefetch = sparse and (not is_distributed)
tmp = helper.create_variable_for_type_inference(dtype)
helper.append_op(
type='lookup_table_v2',
inputs={'Ids': x, 'W': weight},
outputs={'Out': tmp},
attrs={
'is_sparse': sparse,
'is_distributed': is_distributed,
'remote_prefetch': remote_prefetch,
'padding_idx': padding_idx,
},
)
return tmp
python/paddle/nn/functional/loss.py
+5146
View File
File diff suppressed because it is too large Load Diff
python/paddle/nn/functional/moe_permute.py
+177
View File
@@ -0,0 +1,177 @@
# Copyright (c) 2025 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
from typing import TYPE_CHECKING
from paddle import _C_ops
from paddle.base.framework import in_dynamic_or_pir_mode
if TYPE_CHECKING:
from paddle import Tensor
def moe_permute(
hidden_states: Tensor,
scale: Tensor | None,
expert_routemap_topk: Tensor,
expert_prob_topk: Tensor,
num_experts: int,
tokens_per_expert: list,
padding_alignment: int,
do_gather: bool = True,
using_ue8m0_scale: bool = False,
return_expert_indices: bool = False,
override_buffer_size: int = -1,
name: str | None = None,
) -> tuple[Tensor, Tensor, Tensor, Tensor]:
r"""
Permute tokens for Mixture of Experts (MoE) computation in distributed training scenarios.
Note:
This function reorganizes input tokens based on expert assignments to prepare for expert computation.
It handles both bfloat16 and float8_e4m3fn data types with proper scaling for float8 inputs.
1. This function is typically used in pair of moe_unpermute to provide complete MoE functionality.
2. For float8 inputs, proper scaling must be provided via the scale parameter.
3. The padding_alignment parameter affects memory efficiency but not correctness.
4. Any output tokens can find an exact-match in the original input tokens.
5. This permute function has overcomed the aadiff issue, is deterministic.
6. If using_ue8m0_scale is True, then the data type of scale must be int32, and each int32 is packaged from 4 ue8m0 scaling factors.
Args:
hidden_states (Tensor): The input tensor containing tokens to be permuted, stored in row-major layout.
Supported data types: bfloat16 or float8_e4m3fn.
Shape: [sequence_length, token_dimension]
scale (Tensor|None): Scaling factors required when hidden_states is of float8 type.
For float8 inputs, this tensor provides the scaling factors for dequantization.
Shape: [sequence_length, ceil(token_dimension / 128)]. If using_ue8m0_scale is True, the shape is [sequence_length, ceil(ceil(token_dimension / 128)/4)].
Data type: float32 or int32(Only when using_ue8m0_scale is True). If using_ue8m0_scale is True, the data type of scale is int32 which is packed of four ue8m0 scaling factors.
expert_routemap_topk (Tensor): Tensor indicating expert assignments for each token (top-k experts).
Each value represents the expert index the token is assigned to (-1 indicates not assigned).
Shape: [sequence_length, top_k_experts]
Data type: int32
Value range: [-1, num_experts)
expert_prob_topk (Tensor): Tensor containing routing probabilities for top-k experts.
Shape: [sequence_length, top_k_experts]
Data type: float32
num_experts (int): Total number of experts in the MoE layer, limited between 1 and 64.
tokens_per_expert (list[int]): List where each element indicates the number of tokens
assigned to the corresponding expert.
padding_alignment (int): Tokens alignment requirement for expert buffers (in bytes).
Must be a power of 2. Typical values are 16, 32 or 64 for optimal memory access.
do_gather(bool): Decide whether do actual tokens gather operation or not, default is True.
using_ue8m0_scale (bool): Whether to use the ue8m0 scaling for float8 inputs. Default is False.
return_expert_indices(bool): Whether to return an 1D tensor of expert indices for each token, with -1 representing padding. Default is False.
override_buffer_size(bool): Whether to override the buffer size using the given CPU integer, default is -1.
name (str|None, optional): Name prefix for the operation (optional).
Default: None
Returns:
tuple[Tensor, Tensor, Tensor, Tensor]:
- hidden_states_unzipped (Tensor): The permuted and broadcasted input tensor.
Shape: [total_tokens_after_broadcast, token_dimension]
Data type: same as input hidden_states
- zipped_expertwise_rowmap (Tensor): Mapping tensor used to restore original order (unpermute).
Shape: [sequence_length, num_experts]
Data type: int32
- token_prob_unzipped (Tensor): Flattened expert probabilities aligned with permuted tokens.
Shape: [total_tokens_after_broadcast, 1]
Data type: float32
- scale_unzipped (Tensor): Broadcasted scale tensor (only valid for float8 inputs).
Shape: [total_tokens_after_broadcast, scale.shape[-1]]
Data type: float32 or int32. It is same as scale.
Examples:
.. code-block:: pycon
>>> # doctest: +REQUIRES(env:GPU)
>>> # doctest: +SKIP('This is only support in cuda 12.0+')
>>> import paddle
>>> import numpy as np
>>> import paddle.nn.functional as F
>>> hidden_states = paddle.randn([3, 128], dtype='bfloat16')
>>> expert_routemap_topk = paddle.to_tensor(
... [
... [-1, 0, -1, -1, 2, -1, -1, -1],
... [1, -1, -1, -1, -1, -1, -1, -1],
... [-1, -1, -1, -1, -1, -1, 1, -1],
... ],
... dtype='int32',
... )
>>> expert_prob_topk = paddle.to_tensor(
... [
... [0.0, 0.6, 0.0, 0.0, 0.4, 0.0, 0.0, 0.0],
... [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
... [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
... ],
... dtype='float32',
... )
>>> num_experts = 3
>>> tokens_per_expert = [1, 2, 1]
>>> padding_alignment = 2
>>> hidden_states_unzipped, zipped_expertwise_rowmap, token_prob_unzipped, scale_unzipped = F.moe_permute(
... hidden_states,
... None,
... expert_routemap_topk,
... expert_prob_topk,
... num_experts,
... tokens_per_expert,
... padding_alignment,
... )
>>> # weighted by probs.
>>> hidden_states_unzipped = (
... hidden_states_unzipped.astype("float32") * token_prob_unzipped.astype("float32").unsqueeze(-1)
... ).astype("bfloat16")
>>> zipped_tokens, zipped_probs = F.moe_unpermute(
... hidden_states_unzipped, zipped_expertwise_rowmap, expert_routemap_topk, token_prob_unzipped, 3, 3
... )
>>> np.testing.assert_allclose(zipped_tokens.numpy(), hidden_states.numpy(), rtol=1e-05, atol=1e-06)
"""
if in_dynamic_or_pir_mode():
(
hidden_states_unzipped,
zipped_expertwise_rowmap,
token_prob_unzipped,
scale_unzipped,
expert_indices,
) = _C_ops.moe_permute(
hidden_states,
scale,
expert_routemap_topk,
expert_prob_topk,
num_experts,
tokens_per_expert,
padding_alignment,
do_gather,
using_ue8m0_scale,
return_expert_indices,
override_buffer_size,
)
if return_expert_indices:
return (
hidden_states_unzipped,
zipped_expertwise_rowmap,
token_prob_unzipped,
scale_unzipped,
expert_indices,
)
else:
return (
hidden_states_unzipped,
zipped_expertwise_rowmap,
token_prob_unzipped,
scale_unzipped,
)
python/paddle/nn/functional/moe_unpermute.py
+119
View File
@@ -0,0 +1,119 @@
# Copyright (c) 2025 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
from typing import TYPE_CHECKING
from paddle import _C_ops
from paddle.base.framework import in_dynamic_or_pir_mode
if TYPE_CHECKING:
from paddle import Tensor
def moe_unpermute(
hidden_states_unzipped: Tensor,
zipped_expertwise_rowmap: Tensor,
expert_routemap_topk: Tensor,
token_prob_unzipped: Tensor,
total_zipped_tokens: int,
num_experts: int,
using_mix_precision: bool = True,
using_weighted_combine: bool = False,
name: str | None = None,
) -> tuple[Tensor, Tensor]:
r"""
Args:
hidden_states_unzipped (Tensor): The input Tensor containing broadcasted and permuted hidden states.
Shape: (seqlen_broadcasted, token_len). Dtype: bfloat16.
zipped_expertwise_rowmap (Tensor): The input Tensor recording the mapping relationship for unpermute operation.
Shape: (seqlen, num_experts). Dtype: int32.
expert_routemap_topk (Tensor): The input Tensor indicating which expert each token is assigned to.
Shape: (seqlen, 8). Value range: [-1, num_experts]. Dtype: int32.
token_prob_unzipped (Tensor): The input Tensor containing flattened expert probabilities corresponding to hidden_states_unzipped.
Shape: (seqlen_broadcasted, 1). Dtype: float32.
total_zipped_tokens_num (int): The total number of tokens before permutation for output buffer allocation. Dtype: int32.
num_experts (int): The number of experts. Dtype: int32.
using_mix_precision (bool, optional): Whether to use mixed precision during accumulation.
This option significantly improves precision when number of experts > 4. Default: True.
using_weighted_combine (bool, optional): Whether to use weighted token accumulation during unpermute.
Which utilize probs as weights to accumulate tokens. Default: False.
name (str|None, optional): Name for the operation. Default: None.
Returns:
tuple[Tensor, Tensor]: A tuple containing:
- hidden_states (Tensor): The output Tensor with unpermuted tokens.
Shape: (seqlen, token_len). Dtype: bfloat16.
- expert_prob_topk (Tensor): The output Tensor with unpermuted probabilities.
Shape: (seqlen, topk). Dtype: float32.
Examples:
.. code-block:: pycon
>>> # doctest: +REQUIRES(env:GPU)
>>> # doctest: +SKIP('This is only support in cuda 12.0+')
>>> import paddle
>>> import numpy as np
>>> import paddle.nn.functional as F
>>> hidden_states = paddle.randn([3, 128], dtype='bfloat16')
>>> expert_routemap_topk = paddle.to_tensor(
... [
... [-1, 0, -1, -1, 2, -1, -1, -1],
... [1, -1, -1, -1, -1, -1, -1, -1],
... [-1, -1, -1, -1, -1, -1, 1, -1],
... ],
... dtype='int32',
... )
>>> expert_prob_topk = paddle.to_tensor(
... [
... [0.0, 0.6, 0.0, 0.0, 0.4, 0.0, 0.0, 0.0],
... [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
... [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
... ],
... dtype='float32',
... )
>>> num_experts = 3
>>> tokens_per_expert = [1, 2, 1]
>>> padding_alignment = 2
>>> hidden_states_unzipped, zipped_expertwise_rowmap, token_prob_unzipped, scale_unzipped = F.moe_permute(
... hidden_states,
... None,
... expert_routemap_topk,
... expert_prob_topk,
... num_experts,
... tokens_per_expert,
... padding_alignment,
... )
>>> # weighted by probs.
>>> hidden_states_unzipped = (
... hidden_states_unzipped.astype("float32") * token_prob_unzipped.astype("float32").unsqueeze(-1)
... ).astype("bfloat16")
>>> zipped_tokens, zipped_probs = F.moe_unpermute(
... hidden_states_unzipped, zipped_expertwise_rowmap, expert_routemap_topk, token_prob_unzipped, 3, 3
... )
>>> np.testing.assert_allclose(zipped_tokens.numpy(), hidden_states.numpy(), rtol=1e-05, atol=1e-06)
"""
if in_dynamic_or_pir_mode():
zipped_tokens, zipped_probs_topk = _C_ops.moe_unpermute(
hidden_states_unzipped,
zipped_expertwise_rowmap,
expert_routemap_topk,
token_prob_unzipped,
total_zipped_tokens,
num_experts,
using_mix_precision,
using_weighted_combine,
)
return (zipped_tokens, zipped_probs_topk)
python/paddle/nn/functional/norm.py
+921
View File
@@ -0,0 +1,921 @@
# Copyright (c) 2020 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 inspect
import numbers
from typing import TYPE_CHECKING, Any
import numpy as np
from typing_extensions import overload
import paddle
from paddle import _C_ops, in_dynamic_mode
from paddle.base.framework import (
in_dygraph_mode,
in_dynamic_or_pir_mode,
in_pir_mode,
)
from paddle.utils.decorator_utils import (
ParamAliasDecorator,
param_one_alias,
param_two_alias,
)
from ...base.data_feeder import check_type, check_variable_and_dtype
from ...base.layer_helper import LayerHelper
if TYPE_CHECKING:
from collections.abc import Sequence
from typing import Literal
from paddle import Tensor
from paddle._typing import (
DataLayout1D,
DataLayout2D,
DataLayout3D,
DataLayoutND,
)
__all__ = []
@ParamAliasDecorator({"x": ["input"], "axis": ["dim"], "epsilon": ["eps"]})
def normalize(
x: Tensor,
p: float = 2,
axis: int = 1,
epsilon: float = 1e-12,
out: Tensor | None = None,
name: str | None = None,
) -> Tensor:
r"""
Normalize ``x`` along dimension ``axis`` using :math:`L_p` norm. This layer computes
.. math::
y = \frac{x}{ \max\left( \lvert \lvert x \rvert \rvert_p, epsilon\right) }
.. math::
\lvert \lvert x \rvert \rvert_p = \left( \sum_i {\lvert x_i \rvert^p} \right)^{1/p}
where, :math:`\sum_i{\lvert x_i \rvert^p}` is calculated along the ``axis`` dimension.
Parameters:
x (Tensor): The input tensor could be N-D tensor, and the input data type could be float32 or float64.
Alias: ``input``.
p (float|int, optional): The exponent value in the norm formulation. Default: 2.
axis (int, optional): The axis on which to apply normalization. If `axis < 0`, the dimension to normalization is `x.ndim + axis`. -1 is the last dimension.
Alias: ``dim``.
epsilon (float, optional): Small float added to denominator to avoid dividing by zero. Default is 1e-12.
Alias: ``esp``.
name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
out (Tensor|None, optional): The output tensor. Default: None.
Returns:
Tensor, the output has the same shape and data type with ``x``.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn.functional as F
>>> paddle.disable_static()
>>> x = paddle.arange(6, dtype="float32").reshape([2, 3])
>>> y = F.normalize(x)
>>> print(y)
Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[0. , 0.44721359, 0.89442718],
[0.42426407, 0.56568545, 0.70710677]])
>>> y = F.normalize(x, p=1.5)
>>> print(y)
Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[0. , 0.40862012, 0.81724024],
[0.35684019, 0.47578692, 0.59473366]])
>>> y = F.normalize(x, axis=0)
>>> print(y)
Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[0. , 0.24253564, 0.37139067],
[1. , 0.97014254, 0.92847669]])
"""
if in_dygraph_mode():
eps = paddle.full(shape=[1], fill_value=epsilon, dtype=x.dtype)
ret = _C_ops.p_norm(x, float(p), axis, epsilon, True, False)
ret = x / _C_ops.maximum(ret, eps)
elif in_pir_mode():
eps = paddle.full(shape=[1], fill_value=epsilon, dtype=x.dtype)
ret = _C_ops.p_norm(x, float(p), axis, epsilon, True, False)
ret = paddle.divide(x, _C_ops.maximum(ret, eps), name=name)
else:
check_type(p, 'p', (float, int), 'normalize')
check_type(axis, 'axis', (int), 'normalize')
check_variable_and_dtype(
x, 'x', ['float16', 'float32', 'float64'], 'normalize'
)
if len(x.shape) == 1 and axis != 0 and axis != -1:
raise ValueError(
f"Axis must be 0 or -1 when x is a 1-D tensor, but received axis = {axis}"
)
attrs = {
'axis': axis,
'porder': float(p),
'keepdim': True,
'epsilon': epsilon,
}
helper = LayerHelper('p_norm', **locals())
out = helper.create_variable_for_type_inference(dtype=x.dtype)
helper.append_op(
type='p_norm', inputs={'X': x}, outputs={'Out': out}, attrs=attrs
)
eps = out.block.create_var(dtype=out.dtype)
eps = paddle.full(shape=[1], fill_value=epsilon, dtype=out.dtype)
out = paddle.divide(x, paddle.maximum(out, eps), name=name)
if out is not None:
paddle.assign(ret, out)
return out
return ret
@param_two_alias(["x", "input"], ["epsilon", "eps"])
def batch_norm(
x,
running_mean: Tensor,
running_var: Tensor,
weight: Tensor | None = None,
bias: Tensor | None = None,
training: bool = False,
momentum: float = 0.9,
epsilon: float = 1e-05,
data_format: DataLayoutND = 'NCHW',
use_global_stats: bool | None = None,
name: str | None = None,
) -> Tensor:
"""
Applies Batch Normalization as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .
nn.functional.batch_norm is used for nn.BatchNorm1D, nn.BatchNorm2D, nn.BatchNorm3D. Please use above API for BatchNorm.
Parameters:
x(Tensor): input value. It's data type should be float32, float64.
running_mean(Tensor): running mean.
running_var(Tensor): running variance.
weight(Tensor, optional): The weight tensor of batch_norm. Default: None.
bias(Tensor, optional): The bias tensor of batch_norm. Default: None.
epsilon(float, optional): The small value added to the variance to prevent division by zero. Default: 1e-5.
training(bool, optional): True means train mode which compute by batch data and track global mean and var during train period. False means inference mode which compute by global mean and var which calculated by train period. Default False.
momentum(float, optional): The value used for the moving_mean and moving_var computation. Default: 0.9.
data_format(str, optional): Specify the input data format, may be "NC", "NCL", "NCHW", "NCDHW", "NLC", "NHWC" or "NDHWC", where `N` is batch size, `C` is the number of the feature map, `D` is the depth of the feature, `H` is the height of the feature map, `W` is the width of the feature map, `L` is the length of the feature map. Default "NCHW".
use_global_stats(bool|None, optional): Whether to use global mean and variance. If set to False, use the statistics of one mini-batch, if set to True, use the global statistics, if set to None, use global statistics in the test phase and use the statistics of one mini-batch in the training phase. Default: None.
name(str|None, optional): Name for the BatchNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
Returns:
None
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.arange(12, dtype="float32").reshape([2, 1, 2, 3])
>>> print(x)
Tensor(shape=[2, 1, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[0. , 1. , 2. ],
[3. , 4. , 5. ]]],
[[[6. , 7. , 8. ],
[9. , 10., 11.]]]])
>>> running_mean = paddle.to_tensor([0], dtype="float32")
>>> running_variance = paddle.to_tensor([1], dtype="float32")
>>> weight = paddle.to_tensor([2], dtype="float32")
>>> bias = paddle.to_tensor([1], dtype="float32")
>>> batch_norm_out = paddle.nn.functional.batch_norm(
... x,
... running_mean,
... running_variance,
... weight,
... bias,
... )
>>> print(batch_norm_out)
Tensor(shape=[2, 1, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[1. , 2.99998999 , 4.99997997 ],
[6.99996996 , 8.99995995 , 10.99995041]]],
[[[12.99993992, 14.99992943, 16.99991989],
[18.99991035, 20.99990082, 22.99988937]]]])
"""
assert len(x.shape) >= 2, "input dim must be larger than 1"
# input ad out must share the memory
mean_out = running_mean
variance_out = running_var
true_data_format = ['NC', 'NCL', 'NCHW', 'NCDHW', 'NLC', 'NHWC', 'NDHWC']
if data_format not in true_data_format:
raise ValueError(
"data_format must be one of 'NC', 'NCL', 'NCHW', 'NCDHW', "
f"'NLC', 'NHWC', 'NDHWC' but receive {data_format}"
)
data_format = 'NCHW' if data_format[1] == 'C' else 'NHWC'
if use_global_stats is None:
use_global_stats = not training
trainable_statistics = False
else:
trainable_statistics = not use_global_stats
if in_dynamic_mode():
batch_norm_out, _, _, _, _, _ = _C_ops.batch_norm(
x,
running_mean,
running_var,
weight,
bias,
not training,
momentum,
epsilon,
data_format,
use_global_stats,
trainable_statistics,
)
return batch_norm_out
elif in_pir_mode():
batch_norm_out, t1, t2, t3, t4, _ = _C_ops.batch_norm_(
x,
running_mean,
running_var,
weight,
bias,
not training,
momentum,
epsilon,
data_format,
use_global_stats,
trainable_statistics,
)
return batch_norm_out
else:
check_variable_and_dtype(
x, 'input', ['float16', 'uint16', 'float32', 'float64'], 'BatchNorm'
)
# for static need dict
attrs = {
"momentum": momentum,
"epsilon": epsilon,
"is_test": not training,
"data_layout": data_format,
"fuse_with_relu": False,
"use_global_stats": use_global_stats,
"trainable_statistics": trainable_statistics,
}
inputs = {
"X": [x],
"Mean": [running_mean],
"Variance": [running_var],
}
if weight:
inputs['Scale'] = [weight]
if bias:
inputs['Bias'] = [bias]
helper = LayerHelper('batch_norm', **locals())
from paddle.base.data_feeder import convert_dtype
param_dtype = (
x.dtype
if convert_dtype(x.dtype) not in ['float16', 'uint16']
else 'float32'
)
saved_mean = helper.create_variable_for_type_inference(
dtype=param_dtype, stop_gradient=True
)
saved_variance = helper.create_variable_for_type_inference(
dtype=param_dtype, stop_gradient=True
)
batch_norm_out = helper.create_variable_for_type_inference(x.dtype)
outputs = {
"Y": [batch_norm_out],
"MeanOut": [running_mean],
"VarianceOut": [running_var],
"SavedMean": [saved_mean],
"SavedVariance": [saved_variance],
}
if training or trainable_statistics:
# reserve_space is only used for training.
reserve_space = helper.create_variable_for_type_inference(
dtype=x.dtype, stop_gradient=True
)
outputs["ReserveSpace"] = [reserve_space]
helper.append_op(
type="batch_norm", inputs=inputs, outputs=outputs, attrs=attrs
)
return helper.append_activation(batch_norm_out)
@param_two_alias(["x", "input"], ["epsilon", "eps"])
def layer_norm(
x: Tensor,
normalized_shape: int | Sequence[int],
weight: Tensor | None = None,
bias: Tensor | None = None,
epsilon: float = 1e-05,
name: str | None = None,
) -> Tensor:
"""
nn.LayerNorm is recommended.
For more information, please refer to :ref:`api_paddle_nn_LayerNorm` .
.. note::
Alias Support: The parameter name ``input`` can be used as an alias for ``x`` and the parameter name ``eps`` can be used as an alias for ``epsilon``.
For example, ``layer_norm(input=tensor_x, eps=1e-5)`` is equivalent to ``layer_norm(x=tensor_x, epsilon=1e-5)``.
Parameters:
x(Tensor): Input Tensor. It's data type should be bfloat16, float16, float32, float64.
alias: ``input``.
normalized_shape(int|list|tuple): Input shape from an expected input of
size :math:`[*, normalized_shape[0], normalized_shape[1], ..., normalized_shape[-1]]`.
If it is a single integer, this module will normalize over the last dimension
which is expected to be of that specific size.
weight(Tensor, optional): The weight tensor of layer_norm. Default: None.
bias(Tensor, optional): The bias tensor of layer_norm. Default: None.
epsilon(float, optional): The small value added to the variance to prevent
division by zero. Default: 1e-05.
alias: ``eps``.
name(str, optional): Name for the LayerNorm, default is None. For more information, please refer to :ref:`api_guide_Name` .
Returns:
None
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.seed(2023)
>>> x = paddle.rand((2, 2, 2, 3))
>>> layer_norm_out = paddle.nn.functional.layer_norm(x, x.shape[1:])
>>> print(layer_norm_out)
Tensor(shape=[2, 2, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[ 0.87799621, -0.32706589, -1.23529351],
[ 1.01540303, -0.66222930, -0.72354060]],
[[ 1.24183691, 0.45458117, -0.33506936],
[ 0.41468447, 1.26852846, -1.98983335]]],
[[[ 0.02837802, 1.27684653, -0.90110677],
[-0.94709361, -0.15110941, -1.16546965]],
[[-0.82010192, 0.11218391, -0.86506510],
[ 1.09489346, 0.19107464, 2.14656830]]]])
"""
input_shape = list(x.shape)
input_ndim = len(input_shape)
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = [normalized_shape]
elif isinstance(normalized_shape, tuple):
normalized_shape = list(normalized_shape)
elif not isinstance(normalized_shape, list):
raise ValueError(
"`normalized_shape` should be int, list of ints or tuple of ints."
)
normalized_ndim = len(normalized_shape)
begin_norm_axis = input_ndim - normalized_ndim
if input_ndim < normalized_ndim or (
not paddle.utils.is_same_shape(
input_shape[begin_norm_axis:], normalized_shape
)
):
str_normalized_shape = str(normalized_shape)
raise ValueError(
'Given normalized_shape is '
+ str_normalized_shape
+ ', expected input with shape [*, '
+ str_normalized_shape[1:]
+ ', but got input shape '
+ str(input_shape)
)
if in_dynamic_or_pir_mode():
out = _C_ops.layer_norm(x, weight, bias, epsilon, begin_norm_axis)
return out
else:
check_variable_and_dtype(
x, 'input', ['uint16', 'float16', 'float32', 'float64'], 'LayerNorm'
)
inputs = {}
inputs['X'] = [x]
if weight:
inputs['Scale'] = [weight]
if bias:
inputs['Bias'] = [bias]
attrs = {"epsilon": epsilon, "begin_norm_axis": begin_norm_axis}
# create output
helper = LayerHelper('layer_norm', **locals())
from paddle.base.data_feeder import convert_dtype
param_dtype = (
x.dtype if convert_dtype(x.dtype) != 'float16' else 'float32'
)
mean_out = helper.create_variable_for_type_inference(
dtype=param_dtype, stop_gradient=True
)
variance_out = helper.create_variable_for_type_inference(
dtype=param_dtype, stop_gradient=True
)
layer_norm_out = helper.create_variable_for_type_inference(x.dtype)
helper.append_op(
type="layer_norm",
inputs=inputs,
outputs={
"Y": layer_norm_out,
"Mean": mean_out,
"Variance": variance_out,
},
attrs={"epsilon": epsilon, "begin_norm_axis": begin_norm_axis},
)
return helper.append_activation(layer_norm_out)
def rms_norm(
input: Tensor,
normalized_shape: Sequence[int],
weight: Tensor | None = None,
eps: float | None = None,
name: str | None = None,
) -> Tensor:
"""
Applies Layer Normalization over the last dimension of the input tensor using CUDA implementation.
Args:
input (Tensor): Input tensor of shape [rows, cols] or higher dimensions (flattened to 2D).
normalized_shape(list|tuple): Input shape from an expected input of
size :math:`[*, normalized_shape[0], normalized_shape[1], ..., normalized_shape[-1]]`.
If it is a single integer, this module will normalize over the last dimension
which is expected to be of that specific size.
weight(Tensor, optional): The weight tensor of rms_norm. Default: None.
eps(float|None, optional): The small value added to the variance to prevent division by zero.
If None, uses machine epsilon for the compute dtype: ``float64`` inputs use
``np.finfo(np.float64).eps`` (double epsilon), all other dtypes use
``np.finfo(np.float32).eps`` (float epsilon). Default: None.
name (str, optional): Name of the operator.
Returns:
out (Tensor): Normalized tensor of same shape as input.
"""
if eps is None:
if input.dtype == paddle.float64:
eps = np.finfo(np.float64).eps # ~2.22e-16, double machine epsilon
else:
eps = np.finfo(np.float32).eps # ~1.19e-7, float machine epsilon
return _C_ops.rms_norm(input, weight, normalized_shape, eps)[0]
@param_one_alias(["x", "input"])
def instance_norm(
x: Tensor,
running_mean: Tensor | None = None,
running_var: Tensor | None = None,
weight: Tensor | None = None,
bias: Tensor | None = None,
use_input_stats: bool = True,
momentum: float = 0.9,
eps: float = 1e-05,
data_format: Literal['NC', 'NCL', 'NCHW', 'NCDHW'] = 'NCHW',
name: str | None = None,
) -> Tensor:
"""
It is recommended to use :ref:`api_paddle_nn_InstanceNorm1D` , :ref:`api_paddle_nn_InstanceNorm2D` , :ref:`api_paddle_nn_InstanceNorm3D` to call this method internally.
Parameters:
x (Tensor): Input Tensor. It's data type should be float32, float64.
running_mean (Tensor, optional): running mean. Default None. Obsolete (that is, no longer usable).
running_var (Tensor, optional): running variance. Default None. Obsolete (that is, no longer usable).
weight (Tensor, optional): The weight tensor of instance_norm. Default: None.
If its value is None, this parameter will be initialized by one.
bias (Tensor, optional): The bias tensor of instance_norm. Default: None.
If its value is None, this parameter will be initialized by zero.
eps (float, optional): A value added to the denominator for numerical stability. Default is 1e-5.
momentum (float, optional): The value used for the moving_mean and moving_var computation. Default: 0.9.
use_input_stats (bool, optional): Default True. Obsolete (that is, no longer usable).
data_format (str, optional): Specify the input data format, may be "NC", "NCL", "NCHW" or "NCDHW". Default "NCHW".
name (str|None, optional): Name for the InstanceNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
Returns:
None.
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.seed(2023)
>>> x = paddle.rand((2, 2, 2, 3))
>>> instance_norm_out = paddle.nn.functional.instance_norm(x)
>>> print(instance_norm_out)
Tensor(shape=[2, 2, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[ 1.25768495, -0.18054862, -1.26451230],
[ 1.42167914, -0.58056390, -0.65373862]],
[[ 0.95882601, 0.25075224, -0.45947552],
[ 0.21486834, 0.98283297, -1.94780385]]],
[[[ 0.40697321, 1.90885782, -0.71117985],
[-0.76650119, 0.19105314, -1.02920341]],
[[-1.06926346, -0.18710862, -1.11180890],
[ 0.74275863, -0.11246002, 1.73788261]]]])
"""
if in_dynamic_or_pir_mode():
out = _C_ops.instance_norm(x, weight, bias, eps)
return out
else:
check_variable_and_dtype(
x,
'input',
['float32', 'float64', 'float16', 'uint16'],
"InstanceNorm",
)
attrs = {
"epsilon": eps,
"momentum": momentum,
"data_format": data_format,
}
if weight and bias:
inputs = {"X": [x], "Scale": [weight], "Bias": [bias]}
else:
inputs = {"X": [x]}
helper = LayerHelper('instance_norm', **locals())
saved_mean = helper.create_variable_for_type_inference(
dtype=x.dtype, stop_gradient=True
)
saved_variance = helper.create_variable_for_type_inference(
dtype=x.dtype, stop_gradient=True
)
instance_norm_out = helper.create_variable_for_type_inference(x.dtype)
outputs = {
"Y": [instance_norm_out],
"SavedMean": [saved_mean],
"SavedVariance": [saved_variance],
}
helper.append_op(
type="instance_norm", inputs=inputs, outputs=outputs, attrs=attrs
)
return instance_norm_out
def local_response_norm(
x: Tensor,
size: int,
alpha: float = 0.0001,
beta: float = 0.75,
k: float = 1.0,
data_format: DataLayout1D | DataLayout2D | DataLayout3D = 'NCHW',
name: str | None = None,
) -> Tensor:
r"""
Local Response Normalization performs a type of "lateral inhibition" by normalizing over local input regions.
For more information, please refer to `ImageNet Classification with Deep Convolutional Neural Networks <https://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks.pdf>`_
The formula is as follows:
.. math::
Output(i, x, y) = Input(i, x, y) / \left(k + \alpha \sum\limits^{\min(C-1, i + size/2)}_{j = \max(0, i - size/2)}(Input(j, x, y))^2\right)^{\beta}
In the above equation:
- :math:`size` : The number of channels to sum over.
- :math:`k` : The offset (avoid being divided by 0).
- :math:`\\alpha` : The scaling parameter.
- :math:`\\beta` : The exponent parameter.
Args:
x (Tensor): The input 3-D/4-D/5-D tensor. The data type is float16 or float32.
size (int): The number of channels to sum over.
alpha (float, optional): The scaling parameter, positive. Default:1e-4
beta (float, optional): The exponent, positive. Default:0.75
k (float, optional): An offset, positive. Default: 1.0
data_format (str, optional): Specify the data format of the input, and the data format of the output
will be consistent with that of the input. An optional string from:
If x is 3-D Tensor, the string could be `"NCL"` or `"NLC"` . When it is `"NCL"`,
the data is stored in the order of: `[batch_size, input_channels, feature_length]`.
If x is 4-D Tensor, the string could be `"NCHW"`, `"NHWC"`. When it is `"NCHW"`,
the data is stored in the order of: `[batch_size, input_channels, input_height, input_width]`.
If x is 5-D Tensor, the string could be `"NCDHW"`, `"NDHWC"` . When it is `"NCDHW"`,
the data is stored in the order of: `[batch_size, input_channels, input_depth, input_height, input_width]`.
name (str|None, optional): Name for the operation (optional, default is None). For more information,
please refer to :ref:`api_guide_Name`.
Returns:
A tensor storing the transformation result with the same shape and data type as input.
Examples:
.. code-block:: pycon
>>> import paddle
>>> x = paddle.rand(shape=(3, 3, 112, 112), dtype="float32")
>>> y = paddle.nn.functional.local_response_norm(x, size=5)
>>> print(y.shape)
paddle.Size([3, 3, 112, 112])
"""
if not in_dynamic_mode():
check_variable_and_dtype(
x, 'x', ['float16', 'float32'], 'local_response_norm'
)
if data_format not in ['NCL', 'NLC', 'NCHW', 'NHWC', 'NCDHW', 'NDHWC']:
raise ValueError(
"data_format should be in one of [NCL, NCHW, NCDHW, NLC, NHWC, NDHWC], "
f"but got {data_format}"
)
sizes = x.shape
dim = len(sizes)
if dim < 3:
raise ValueError(
f'Expected 3D or higher dimensionality input, but got {dim} dimensions'
)
for i, sz in enumerate(sizes):
if not sz > 0 and i > 0:
raise ValueError(
"Expected every dim's size to be larger than 0, "
f"but the size of the {i}-th dim is {sz}"
)
channel_last = True if data_format[-1] == "C" else False
from functools import reduce
sum_sizes = reduce(lambda x, y: x * y, sizes[1:], 1)
div = paddle.unsqueeze(paddle.multiply(x, x), axis=1)
if not channel_last:
pad4d_shape = [0, 0, size // 2, (size - 1) // 2]
pool2d_shape = (size, 1)
reshape_shape = [
sizes[0],
1,
sizes[1],
sizes[2],
int(sum_sizes / (sizes[1] * sizes[2])),
]
pad5d_shape = [0, 0, 0, 0, size // 2, (size - 1) // 2]
pool3d_shape = (size, 1, 1)
else:
pad4d_shape = [size // 2, (size - 1) // 2, 0, 0]
pool2d_shape = (1, size)
reshape_shape = [
sizes[0],
1,
sizes[1],
int(sum_sizes / (sizes[1] * sizes[-1])),
sizes[-1],
]
pad5d_shape = [size // 2, (size - 1) // 2, 0, 0, 0, 0]
pool3d_shape = (1, 1, size)
if dim == 3:
div = paddle.nn.functional.pad(div, pad=pad4d_shape)
div = paddle.nn.functional.avg_pool2d(
div, kernel_size=pool2d_shape, stride=1
)
div = paddle.squeeze(div, axis=1)
else:
div = paddle.reshape(div, shape=reshape_shape)
div = paddle.nn.functional.pad(
div, pad=pad5d_shape, data_format='NCDHW'
)
div = paddle.nn.functional.avg_pool3d(
div, kernel_size=pool3d_shape, stride=1
)
div = paddle.reshape(paddle.squeeze(div, axis=1), sizes)
div = paddle.scale(div, scale=alpha, bias=k)
div = paddle.pow(div, beta)
res = paddle.divide(x, div, name=name)
return res
@overload
def group_norm(
x: Tensor,
num_groups: int,
epsilon: float = 1e-05,
weight: Tensor | None = None,
bias: Tensor | None = None,
data_format: DataLayout1D | DataLayout2D | DataLayout3D = 'NCHW',
name: str | None = None,
) -> Tensor: ...
@overload
def group_norm(
input: Tensor,
num_groups: int,
weight: Tensor | None = None,
bias: Tensor | None = None,
eps: float = 1e-05,
) -> Tensor: ...
def group_norm(*args: Any, **kwargs: Any) -> Tensor:
"""
nn.GroupNorm is recommended.
For more information, please refer to :ref:`api_paddle_nn_GroupNorm` .
This function has two functionalities, depending on the parameters passed:
1. ``group_norm(Tensor input, int num_groups, Tensor weight = None, Tensor bias = None, float eps = 1e-05)``:
PyTorch compatible group_norm.
2. ``group_norm(Tensor x, int num_groups, float epsilon = 1e-05, Tensor weight = None, Tensor bias = None,
DataLayout1D | DataLayout2D | DataLayout3D data_format = 'NCHW', str | None name = None)``:
The original paddle.nn.functional.group_norm, see the following docs.
Parameters:
x(Tensor): Input Tensor with shape: attr:`(batch, num_features, *)`.
alias: ``input``.
num_groups(int): The number of groups that divided from channels.
epsilon(float, optional): The small value added to the variance to prevent
division by zero. Default: 1e-05.
alias: ``eps``.
weight(Tensor, optional): The weight Tensor of group_norm, with shape: attr:`[num_channels]`.
Default: None.
bias(Tensor, optional): The bias Tensor of group_norm, with shape: attr:`[num_channels]`.
Default: None.
data_format(str, optional): Specify the input data format. Support "NCL", "NCHW", "NCDHW", "NLC", "NHWC" or "NDHWC". Default: "NCHW".
name(str|None, optional): Name for the GroupNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
Returns:
Tensor, the output has the same shape with ``x``.
Examples:
.. code-block:: pycon
>>> import paddle
>>> paddle.seed(100)
>>> x = paddle.arange(48, dtype="float32").reshape((2, 6, 2, 2))
>>> group_norm_out = paddle.nn.functional.group_norm(x, num_groups=6)
>>> print(group_norm_out)
Tensor(shape=[2, 6, 2, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]]],
[[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]],
[[-1.34163547, -0.44721183],
[ 0.44721183, 1.34163547]]]])
"""
len_args = len(args)
if len_args + len(kwargs) < 2:
raise TypeError(
f"Too few arguments in the function call: {len_args}, {len(kwargs)}. Expect one of: \n"
" - (Tensor input, int num_groups, Tensor weight = None, Tensor bias = None, float eps = 1e-05)\n"
" - (Tensor x, int num_groups, float epsilon = 1e-05, Tensor weight = None, Tensor bias = None, "
"DataLayout1D | DataLayout2D | DataLayout3D data_format = 'NCHW', str | None name = None)"
)
def safe_set_param(key: str, value: Any):
if key in kwargs:
raise TypeError(f"got multiple values for argument '{key}'")
kwargs[key] = value
if 'input' in kwargs:
safe_set_param('x', kwargs.pop('input'))
if 'eps' in kwargs:
safe_set_param('epsilon', kwargs.pop('eps'))
if len_args >= 3 and not isinstance(args[2], float):
param_keys = ["weight", "bias", "epsilon"]
for idx in range(min(len_args - 2, len(param_keys))):
safe_set_param(param_keys[idx], args[idx + 2])
args = args[:2]
return _group_norm_wrapper(*args, **kwargs)
def _group_norm_wrapper(
x: Tensor,
num_groups: int,
epsilon: float = 1e-05,
weight: Tensor | None = None,
bias: Tensor | None = None,
data_format: DataLayout1D | DataLayout2D | DataLayout3D = 'NCHW',
name: str | None = None,
) -> Tensor:
if data_format not in ['NCL', 'NCHW', 'NCDHW', 'NLC', 'NHWC', 'NDHWC']:
raise ValueError("unsupported data layout:" + data_format)
data_format = 'NCHW' if data_format[1] == 'C' else 'NHWC'
if in_dynamic_or_pir_mode():
return _C_ops.group_norm(
x,
weight,
bias,
epsilon,
num_groups,
data_format,
)
else:
helper = LayerHelper('group_norm', **locals())
mean_out = helper.create_variable_for_type_inference(
dtype=x.dtype, stop_gradient=True
)
variance_out = helper.create_variable_for_type_inference(
dtype=x.dtype, stop_gradient=True
)
inputs = {'X': x}
if bias is not None:
inputs['Bias'] = bias
if weight is not None:
inputs['Scale'] = weight
# create output
group_norm_out = helper.create_variable_for_type_inference(
dtype=x.dtype
)
helper.append_op(
type="group_norm",
inputs=inputs,
outputs={
"Y": group_norm_out,
"Mean": mean_out,
"Variance": variance_out,
},
attrs={
"epsilon": epsilon,
"groups": num_groups,
"data_layout": data_format,
},
)
return helper.append_activation(group_norm_out)
group_norm.__signature__ = inspect.signature(_group_norm_wrapper)
python/paddle/nn/functional/pooling.py
Executable
+2822
View File
File diff suppressed because it is too large Load Diff
python/paddle/nn/functional/sdpa.py
+681
View File
@@ -0,0 +1,681 @@
# Copyright (c) 2025 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
from dataclasses import dataclass
from functools import cached_property, lru_cache
from typing import TYPE_CHECKING
import paddle
from paddle import _C_ops
from paddle.base.log_helper import get_logger
from paddle.nn.attention.sdpa import (
SDPBackend,
_get_backend_priority,
_get_enabled_backends,
)
_logger = get_logger(
__name__, "INFO", fmt='%(asctime)s-%(levelname)s: %(message)s'
)
from paddle.nn.functional.flash_attention import _math_attention
if TYPE_CHECKING:
from paddle import Tensor, dtype
from paddle.base.core import Place
_config = {}
def init_config():
global _config
_config = {
"flash_attn": {
"MINIMUM_SM_VERSION": (8, 0),
"MAXIMUM_SM_VERSION": (12, 1),
"support_dtypes": (paddle.float16, paddle.bfloat16)
if paddle.device.is_bf16_supported(including_emulation=False)
else (paddle.float16,),
},
"mem_efficient_attn": {
"MINIMUM_SM_VERSION": (5, 0),
"MAXIMUM_SM_VERSION": (12, 1),
"support_dtypes": (
paddle.float16,
paddle.bfloat16,
paddle.float,
)
if paddle.device.is_bf16_supported(including_emulation=False)
else (paddle.float16, paddle.float),
},
}
def _repeat_kv(key: Tensor, value: Tensor, num_repeats: int):
"""
Repeat key and value tensors along the num_heads(3) dimension. The layout
of key and value should be [batch_size, seq_len, num_heads, head_dim].
"""
if num_repeats == 1:
return key, value
# repeat_interleave does not support float16 on GPU, so we manually expand the tensor
key, value = key.unsqueeze(3), value.unsqueeze(3)
key, value = (
key.expand([-1, -1, -1, num_repeats, -1]),
value.expand([-1, -1, -1, num_repeats, -1]),
)
key, value = (
key.flatten(2, 3).contiguous(),
value.flatten(2, 3).contiguous(),
)
return key, value
@dataclass
class SDPParams:
query_shape: paddle.Size
key_shape: paddle.Size
value_shape: paddle.Size
attn_mask_shape: paddle.Size | None
dropout: float
is_causal: bool
scale: float | None
query_stop_gradient: bool
dtype: tuple[dtype, dtype, dtype]
place: tuple[Place, Place, Place]
@cached_property
def batch_size(self) -> tuple[int, int, int]:
return self.query_shape[0], self.key_shape[0], self.value_shape[0]
@cached_property
def seq_len(self) -> tuple[int, int, int]:
return self.query_shape[1], self.key_shape[1], self.value_shape[1]
@cached_property
def num_heads(self) -> tuple[int, int, int]:
return self.query_shape[2], self.key_shape[2], self.value_shape[2]
@cached_property
def head_dim(self) -> tuple[int, int, int]:
return self.query_shape[-1], self.key_shape[-1], self.value_shape[-1]
@cached_property
def device_id(self) -> tuple[int, ...]:
ret = tuple(
pl.gpu_device_id() if pl.is_gpu_place() else -1 for pl in self.place
)
return ret
@lru_cache(maxsize=8)
def get_device_capability(device_id: int) -> tuple[int, int]:
if device_id < 0:
return (0, 0)
return paddle.device.cuda.get_device_capability(device_id)
@lru_cache(maxsize=32)
def check_sm_version(
min_sm: tuple[int, int], max_sm: tuple[int, int], device_id: int = 0
) -> bool:
major, minor = get_device_capability(device_id)
current = (major, minor)
return min_sm <= current <= max_sm
@lru_cache(maxsize=1)
def check_cuda_is_available() -> bool:
return paddle.is_compiled_with_cuda() and paddle.cuda.is_available()
def check_all_tensors_on_device(params: SDPParams):
"""
Check all input tensors are placed on the GPU device.
"""
if not (
params.place[0].is_gpu_place() or params.place[0].is_custom_place()
):
_logger.debug(
"All input tensors should be placed on GPU or custom place, but "
f"query place: {params.place[0]}, key place: "
f"{params.place[1]}, value place: {params.place[2]}"
)
return False
return True
def check_head_dim_size_flash(params: SDPParams):
"""
Check the dimension of head in query, key, and value should be equal and all less than 256.
"""
q_head_dim, k_head_dim, v_head_dim = params.head_dim
if q_head_dim > 256 or q_head_dim != k_head_dim or k_head_dim != v_head_dim:
_logger.debug(
"The dimension of head in query, key, and value should be equal and all less than 256, "
f"but q_head_dim: {q_head_dim}, k_head_dim: {k_head_dim}, v_head_dim: {v_head_dim}"
)
return False
if q_head_dim % 8 != 0:
_logger.debug(
"The dimension of head in query, key, and value should be a multiple of 8, "
f"but q_head_dim: {q_head_dim}"
)
return False
return True
@lru_cache(maxsize=8)
def check_flash_attention_hardware_support(device_id: int):
"""
Check flash attention requires CUDA support and SM between 8.0 and 12.1.
"""
if SDPBackend.FLASH_ATTENTION and paddle.is_compiled_with_custom_device(
paddle.device.get_all_device_type()[0]
):
return True
if not check_cuda_is_available():
_logger.debug("Flash attention requires CUDA support.")
return False
if not check_sm_version(
_config["flash_attn"]["MINIMUM_SM_VERSION"],
_config["flash_attn"]["MAXIMUM_SM_VERSION"],
device_id,
):
_logger.debug(
f"Flash attention requires SM between {_config['flash_attn']['MINIMUM_SM_VERSION']}"
f"and {_config['flash_attn']['MAXIMUM_SM_VERSION']}, but found SM "
f"{get_device_capability(device_id)}"
)
return False
return True
def check_flash_causal_non_square_seqlens(params: SDPParams):
"""
Check flash attention only supports causal attention when the sequence length of query and key are equal.
"""
if not params.is_causal:
return True
q_len, k_len, _ = params.seq_len
if q_len == k_len:
return True
_logger.debug(
f"Flash attention only supports causal attention when the sequence"
f"length of query and key are equal, but got query shape: "
f"{params.query_shape}, key shape: {params.key_shape}"
)
return False
def check_dtypes_low_precision_fa(params: SDPParams):
"""
check QKV share the same dtype and are supported dtype.
"""
q_dtype, k_dtype, v_dtype = params.dtype
if (
q_dtype != k_dtype
or v_dtype != k_dtype
or q_dtype not in _config["flash_attn"]["support_dtypes"]
):
_logger.debug(
f"Flash attention requires query, key, and value "
f"to be of the same dtype and support dtype, but "
f"got query dtype: {q_dtype}, key dtype: {k_dtype}"
f", value dtype: {v_dtype}. Supported dtypes are: "
f"{_config['flash_attn']['support_dtypes']}"
)
return False
return True
def check_dtypes_low_precision_mem_efficient_attn(params: SDPParams):
"""
check QKV share the same dtype and are supported dtype.
"""
q_dtype, k_dtype, v_dtype = params.dtype
if (
q_dtype != k_dtype
or v_dtype != k_dtype
or q_dtype not in _config["mem_efficient_attn"]["support_dtypes"]
):
_logger.debug(
f"Mem_efficient_attn requires query, key, and value "
f"to be of the same dtype and support dtype, but "
f"got query dtype: {q_dtype}, key dtype: {k_dtype}"
f", value dtype: {v_dtype}. Supported dtypes are: "
f"{_config['mem_efficient_attn']['support_dtypes']}"
)
return False
return True
@lru_cache(maxsize=2)
def use_tensor_cores(is_half: bool, device_id: int) -> bool:
major, _ = get_device_capability(device_id)
if major >= 8:
return True
if major == 7:
return is_half
return False
@lru_cache(maxsize=32)
def minimum_gemm_alignment(dtype: dtype, device_id: int):
is_half = dtype in (paddle.float16, paddle.bfloat16)
use_tc = use_tensor_cores(is_half, device_id)
major, _ = get_device_capability(device_id)
matmul_alignment_mn = 4 if major > 8 else 1
bits_per_scalar = 16 if is_half else 32
if use_tc:
matmul_alignment_mn = max(matmul_alignment_mn, 128 / bits_per_scalar)
return matmul_alignment_mn
@lru_cache(maxsize=8)
def check_mem_efficient_hardware_support(device_id: int):
"""
Check mem_efficient attention requires CUDA support and SM between 5.0 and 12.1.
"""
if not check_cuda_is_available():
_logger.debug("Mem efficient attention requires CUDA support.")
return False
if not check_sm_version(
_config["mem_efficient_attn"]["MINIMUM_SM_VERSION"],
_config["mem_efficient_attn"]["MAXIMUM_SM_VERSION"],
device_id,
):
_logger.debug(
f"Mem efficient attention requires SM between {_config['mem_efficient_attn']['MINIMUM_SM_VERSION']}"
f"and {_config['mem_efficient_attn']['MAXIMUM_SM_VERSION']}, but found SM "
f"{get_device_capability(device_id)}"
)
return False
return True
def check_head_dim_size_mem_efficient(params: SDPParams):
q_head_dim, k_head_dim, v_head_dim = (
params.query_shape[-1],
params.key_shape[-1],
params.value_shape[-1],
)
alignment = minimum_gemm_alignment(params.dtype[0], params.device_id[0])
if (
q_head_dim % alignment != 0
or k_head_dim % alignment != 0
or v_head_dim % alignment != 0
):
_logger.debug(
f"Mem efficient attention requires head dim size aligned to {alignment}, "
f"but found q_head_dim: {q_head_dim}, k_head_dim: {k_head_dim}, v_head_dim: {v_head_dim}"
)
return False
return True
def check_attn_mask_alignment(params: SDPParams) -> bool:
if params.is_causal:
return True
if params.attn_mask_shape is None:
return True
last_dim = params.attn_mask_shape[-1]
if last_dim % 8 != 0:
_logger.debug(
f"Mem efficient attention requires attn_mask last dimension to be divisible by 8 "
f"to satisfy vector alignment, but got {last_dim}. "
"Falling back to other backends."
)
return False
return True
def check_scale_is_None(params: SDPParams) -> bool:
if params.scale is None:
return True
_logger.debug("Paddle's FAV2 does not support scale parameter.")
return False
def can_use_flash_attention(params: SDPParams = False) -> bool:
general_constraints = [
check_all_tensors_on_device,
check_head_dim_size_flash,
check_flash_causal_non_square_seqlens,
check_dtypes_low_precision_fa,
check_scale_is_None,
]
for constraint in general_constraints:
if not constraint(params):
return False
if not check_flash_attention_hardware_support(params.device_id[0]):
return False
return True
def can_use_mem_efficient_attention(params: SDPParams = False) -> bool:
constraints = [
check_all_tensors_on_device,
check_head_dim_size_mem_efficient,
check_attn_mask_alignment,
check_dtypes_low_precision_mem_efficient_attn,
]
for constraint in constraints:
if not constraint(params):
return False
if not check_mem_efficient_hardware_support(params.device_id[0]):
return False
return True
def select_sdp_for_sdpa(param: SDPParams) -> str:
# Note: This API is designed for nn.functional.scaled_dot_product_attention,
# and is **NOT** expected to be called by others. Some promises should be guaranteed
# by caller to skip some rarely unmet constraints:
# 1. The input dim is 4, layout is (batch, seq_len, num_heads, head_dim)
# 2. The batch_size and num_heads of each input should be the same
place = paddle.get_device()
if "xpu" in place:
return "flash_attn"
enabled_backends = _get_enabled_backends()
priority_order = _get_backend_priority()
for backend in priority_order:
if backend not in enabled_backends:
continue
if backend == SDPBackend.FLASH_ATTENTION:
if can_use_flash_attention(param):
return "flash_attn"
elif backend == SDPBackend.EFFICIENT_ATTENTION:
if can_use_mem_efficient_attention(param):
return "mem_efficient"
elif backend == SDPBackend.MATH:
return "math"
raise RuntimeError(
"No available backend for scaled_dot_product_attention was found."
)
def scaled_dot_product_attention(
query: Tensor,
key: Tensor,
value: Tensor,
attn_mask: Tensor | None = None,
dropout_p: float = 0.0,
is_causal: bool = False,
training: bool = True,
backend: str | None = None,
scale: float | None = None,
enable_gqa: bool = True,
name: str | None = None,
) -> Tensor:
r"""
The equation is:
.. math::
result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V
where : ``Q``, ``K``, and ``V`` represent the three input parameters of the attention module.
The dimensions of the three parameters are the same.
``d`` represents the size of the last dimension of the three parameters.
Warning:
This API only verifies inputs with dtype float16 and bfloat16, other dtypes may fall back to math
implementation, which is less optimized.
Warning:
If is_causal is set to True, the causal mask should not be provided, otherwise
the provided mask will be ignored.
Note:
This API differs from :ref:`api_paddle_compat_nn_functional_scaled_dot_product_attention` in that:
1. The QKV layout of this API is [batch_size, seq_len, num_heads, head_dim] or [seq_len, num_heads, head_dim].
If you need num_heads before seq_len layout, please use ``paddle.compat.nn.functional.scaled_dot_product_attention``.
Args:
query(Tensor): The query tensor in the Attention module.
4-D tensor with shape:
[batch_size, seq_len_key, num_heads, head_dim].
3-D tensor with shape:
[seq_len_key, num_heads, head_dim].
The dtype can be float16 or bfloat16.
key(Tensor): The key tensor in the Attention module.
4-D tensor with shape:
[batch_size, seq_len_key, num_heads, head_dim].
3-D tensor with shape:
[seq_len_key, num_heads, head_dim].
The dtype can be float16 or bfloat16.
value(Tensor): The value tensor in the Attention module.
4-D tensor with shape:
[batch_size, seq_len_value, num_heads, head_dim].
3-D tensor with shape:
[seq_len_value, num_heads, head_dim].
The dtype can be float16 or bfloat16.
attn_mask(Tensor, optional): The attention mask tensor. The shape should be broadcastable to
[batch_size, num_heads, seq_len_key, seq_len_query]. The dtype can be bool
or same type of query. The bool mask indicates the positions should take part
in attention. The non-bool mask will be added to attention score.
dropout_p(float, optional): The dropout ratio.
is_causal(bool, optional): Whether enable causal mode.
training(bool, optional): Whether it is in the training phase.
backend(str, optional): Specify which backend to compute scaled dot product attention.
Currently only support "p2p" for distribution usage.
scale(float, optional): The scaling factor used in the calculation of attention weights.
If None, scale = 1 / sqrt(head_dim).
enable_gqa(bool, optional): Whether enable GQA(Group Query Attention) mode. Default is True.
name(str|None, optional): The default value is None. Normally there is no need for user
to set this property. For more information, please refer to
:ref:`api_guide_Name`.
Returns:
out(Tensor): The attention tensor.
4-D tensor with shape: [batch_size, seq_len, num_heads, head_dim].
3-D tensor with shape: [seq_len, num_heads, head_dim].
The dtype can be float16 or bfloat16.
Examples:
.. code-block:: pycon
>>> # doctest: +SKIP('bfloat need V100 compile')
>>> import paddle
>>> q = paddle.rand((1, 128, 2, 16), dtype=paddle.bfloat16)
>>> output = paddle.nn.functional.scaled_dot_product_attention(q, q, q, None, 0.9, False)
>>> print(output)
>>> # doctest: -SKIP
"""
is_batched = query.dim() == 4
if not is_batched:
# FlashAttention backend does not support unbatched input,
# we add batch dim here and will skip check input dim when selecting FA backend.
query = query.unsqueeze(0)
key = key.unsqueeze(0)
value = value.unsqueeze(0)
k_heads, q_heads, v_heads = (
key.shape[2],
query.shape[2],
value.shape[2],
)
if enable_gqa:
assert k_heads == 0 or q_heads % k_heads == 0, (
f"The number of groups in query({q_heads}) must be divisible by the number of groups in key({k_heads}) if GQA enabled."
)
assert k_heads == v_heads, (
f"The number of groups in key({k_heads}) must be equal to the number of groups in value({v_heads}) if GQA enabled."
)
else:
assert q_heads == k_heads == v_heads, (
f"The number of groups in query({q_heads}) must be equal to the number of groups in key({k_heads}) "
f"and the number of groups in value({v_heads}) if GQA disabled."
)
bs, seq_len_q, num_heads_q, head_dim_q = query.shape
_, seq_len_k, num_heads_k, head_dim_k = key.shape
if (
backend == 'p2p'
and query.is_dist()
and key.is_dist()
and value.is_dist()
):
# ring attention for auto_parallel mode
assert scale is None, f"Backend {backend} not support scale parameter."
out = paddle.distributed.auto_parallel.ring_attention.RingFlashAttention.apply(
query,
key,
value,
attn_mask,
dropout_p,
is_causal,
)
return out
if not paddle.base.in_dygraph_mode():
qkv_place = (paddle.framework._current_expected_place_(),) * 3
else:
qkv_place = (query.place, key.place, value.place)
param = SDPParams(
query_shape=query.shape,
key_shape=key.shape,
value_shape=value.shape,
attn_mask_shape=attn_mask.shape if attn_mask is not None else None,
dropout=dropout_p,
is_causal=is_causal,
scale=scale,
query_stop_gradient=query.stop_gradient,
dtype=(query.dtype, key.dtype, value.dtype),
place=qkv_place,
)
if len(_config) == 0:
init_config()
is_zero_size = (
query.shape.numel() == 0
or key.shape.numel() == 0
or value.shape.numel() == 0
)
if attn_mask is not None:
if attn_mask.dtype == paddle.bool:
attn_mask = paddle.where(
attn_mask,
paddle.to_tensor(0.0, dtype=query.dtype),
paddle.to_tensor(-float('inf'), dtype=query.dtype),
)
if is_zero_size:
sdp_func_name = "math"
else:
sdp_func_name = select_sdp_for_sdpa(param)
_logger.debug("Selected backend:" + sdp_func_name)
if sdp_func_name == "flash_attn":
fixed_seed_offset = None
return_softmax = False
rng_name = ""
if attn_mask is not None:
if attn_mask.ndim == 2:
attn_mask = attn_mask.expand([bs, 1, *attn_mask.shape])
elif attn_mask.ndim == 3:
attn_mask = paddle.unsqueeze(attn_mask, axis=1)
out, _, _, _ = _C_ops.flash_attn(
query,
key,
value,
fixed_seed_offset,
attn_mask,
dropout_p,
is_causal,
return_softmax,
not training,
rng_name,
)
elif sdp_func_name == "mem_efficient":
from paddle.incubate.nn.memory_efficient_attention import (
LowerTriangularMask,
memory_efficient_attention,
)
repeats = q_heads // k_heads
key, value = _repeat_kv(key, value, repeats)
if is_causal:
attn_mask = LowerTriangularMask()
elif attn_mask is not None:
# if need broadcast, memory_efficient_attention requires to
# broadcast first two dim simultaneously
if attn_mask.dim() == 3:
attn_mask = attn_mask.unsqueeze(axis=1)
if attn_mask.dim() == 4 and (
attn_mask.shape[0] != bs ^ attn_mask.shape[1] != num_heads_q
):
attn_mask = attn_mask.expand(
[
bs,
num_heads_q,
attn_mask.shape[2],
attn_mask.shape[3],
]
)
out = memory_efficient_attention(
query,
key,
value,
attn_bias=attn_mask,
p=dropout_p,
scale=scale,
training=training,
)
elif sdp_func_name == "math":
repeats = q_heads // k_heads if k_heads != 0 else 1
key, value = _repeat_kv(key, value, repeats)
if attn_mask is not None and attn_mask.dim() == 3:
attn_mask = attn_mask.unsqueeze(axis=1)
out = _math_attention(
query,
key,
value,
attn_mask,
dropout_p,
is_causal,
False,
training,
scale,
)[0]
else:
raise ValueError(f"Invalid backend {backend}")
if not is_batched:
out = paddle.squeeze(out, axis=0)
return out
python/paddle/nn/functional/sparse_attention.py
+179
View File
@@ -0,0 +1,179 @@
# 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.
from __future__ import annotations
import paddle
from paddle import _C_ops
from paddle.base.layer_helper import LayerHelper
def sparse_attention(
query: paddle.Tensor,
key: paddle.Tensor,
value: paddle.Tensor,
sparse_csr_offset: paddle.Tensor,
sparse_csr_columns: paddle.Tensor,
key_padding_mask: paddle.Tensor | None = None,
attn_mask: paddle.Tensor | None = None,
name: str | None = None,
) -> paddle.Tensor:
r"""
This operator sparsify the Attention matrix in Transformer module
to achieve the effect of reducing memory consumption and computation.
The sparse layout is expressed in CSR format and contains two parameters,
``offset`` and ``columns``. The equation is:
.. math::
result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V
where : ``Q``, ``K``, and ``V`` represent the three input parameters of the attention module.
The dimensions of the three parameters are the same.
``d`` represents the size of the last dimension of the three parameters.
Warning:
This API is only used in ``CUDA 11.3`` and above versions.
Args:
query(Tensor): The query tensor in the Attention module.
4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
The dtype can be float32 and float64.
key(Tensor): The key tensor in the Attention module.
4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
The dtype can be float32 and float64.
value(Tensor): The value tensor in the Attention module.
4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
The dtype can be float32 and float64.
sparse_csr_offset(Tensor): The sparsity feature in the Attention module
is expressed in the CSR format, and the offset represents
the number of non-zero elements in each row of the matrix.
3-D tensor with shape:
[batch_size, num_heads, seq_len + 1].
The dtype should be int32.
sparse_csr_columns(Tensor): The sparsity feature in the Attention module
is expressed in the CSR format, and the columns represent
the column index values of non-zero elements in the matrix.
3-D tensor with shape:
[batch_size, num_heads, sparse_nnz].
The dtype should be int32.
key_padding_mask(Tensor|None, optional):The key padding mask tensor in the Attention module.
2-D tensor with shape: [batch_size, seq_len].
The dtype can be float32 and float64.
A value of 0 means that the position is masked.
attn_mask(Tensor|None, optional):The attention mask tensor in the Attention module.
2-D tensor with shape: [seq_len, seq_len].
The dtype can be float32 and float64.
A value of 0 means that the position is masked.
name(str|None, optional): The default value is None. Normally there is no need for user
to set this property. For more information, please refer to
:ref:`api_guide_Name`.
Returns:
Tensor, 4-D tensor with shape:
[batch_size, num_heads, seq_len, head_dim].
The dtype can be float32 or float64.
Examples:
.. code-block:: pycon
>>> # doctest: +SKIP('This API is only used in CUDA11.3 and above.')
>>> import paddle
>>> paddle.disable_static()
>>> # `query`, `key` and `value` all have shape [1, 1, 4, 2]
>>> query = paddle.to_tensor(
... [[[[0, 1], [2, 3], [0, 1], [2, 3]]]],
... dtype="float32",
... )
>>> key = paddle.to_tensor([[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32")
>>> value = paddle.to_tensor([[[[0, 1], [2, 3], [0, 1], [2, 3]]]], dtype="float32")
>>> offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
>>> columns = paddle.to_tensor([[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32")
>>> print(offset.shape)
paddle.Size([1, 1, 5])
>>> print(columns.shape)
paddle.Size([1, 1, 8])
...
>>> key_padding_mask = paddle.to_tensor([[1, 1, 1, 0]], dtype="float32")
>>> attention_mask = paddle.to_tensor(
... [
... [1, 0, 1, 1],
... [1, 1, 1, 1],
... [1, 1, 1, 1],
... [1, 1, 1, 1],
... ],
... dtype="float32",
... )
>>> output_mask = paddle.nn.functional.sparse_attention(
... query,
... key,
... value,
... offset,
... columns,
... key_padding_mask=key_padding_mask,
... attn_mask=attention_mask,
... )
>>> print(output_mask)
Tensor(shape=[1, 1, 4, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
[[[[0. , 1. ],
[1.99830270, 2.99830270],
[0. , 1. ],
[0. , 1. ]]]])
>>> output = paddle.nn.functional.sparse_attention(query, key, value, offset, columns)
>>> print(output)
Tensor(shape=[1, 1, 4, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
[[[[1.60885942, 2.60885954],
[1.99830270, 2.99830270],
[1.60885942, 2.60885954],
[1.99830270, 2.99830270]]]])
"""
if paddle.framework.in_dynamic_or_pir_mode():
res = _C_ops.sparse_attention(
query,
key,
value,
sparse_csr_offset,
sparse_csr_columns,
key_padding_mask,
attn_mask,
)
return res
helper = LayerHelper('sparse_attention', **locals())
dtype = helper.input_dtype(input_param_name='Q')
out = helper.create_variable_for_type_inference(dtype)
result_sdd = helper.create_variable_for_type_inference(dtype)
result_softmax = helper.create_variable_for_type_inference(dtype)
inputs = {
'Q': query,
'K': key,
'V': value,
'Offset': sparse_csr_offset,
'Columns': sparse_csr_columns,
'KeyPaddingMask': key_padding_mask,
'AttnMask': attn_mask,
}
outputs = {
'Out': out,
'SparseDotSdd': result_sdd,
'Softmax': result_softmax,
}
helper.append_op(type='sparse_attention', inputs=inputs, outputs=outputs)
return out
python/paddle/nn/functional/transformer.py
+16
View File
@@ -0,0 +1,16 @@
# Copyright (c) 2020 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.
# TODO: define the classes of Transformer neural network
# __all__ = [ ]
python/paddle/nn/functional/vision.py
+300
View File
@@ -0,0 +1,300 @@
# Copyright (c) 2020 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
from typing import TYPE_CHECKING
import paddle
from paddle import _C_ops, in_dynamic_mode
from paddle._C_ops import (
grid_sample, # noqa: F401
pixel_shuffle, # noqa: F401
)
from paddle.base.framework import (
in_dynamic_or_pir_mode,
in_pir_mode,
)
from paddle.utils.decorator_utils import param_one_alias
from ...base.data_feeder import check_variable_and_dtype
from ...base.layer_helper import LayerHelper
from ...common_ops_import import Variable
from ...device import get_cudnn_version, is_compiled_with_rocm
if TYPE_CHECKING:
from paddle import Tensor
from paddle._typing import DataLayout2D, ShapeLike
__all__ = []
def affine_grid(
theta: Tensor,
out_shape: ShapeLike,
align_corners: bool = True,
name: str | None = None,
) -> Tensor:
"""
It generates a grid of (x,y) or (x,y,z) coordinates using the parameters of
the affine transformation that correspond to a set of points where
the input feature map should be sampled to produce the transformed
output feature map.
Args:
theta (Tensor): A tensor with shape [N, 2, 3] or [N, 3, 4]. It contains a batch of affine transform parameters.
The data type can be float32 or float64.
out_shape (Tensor | list | tuple): Type can be a 1-D Tensor, list, or tuple. It is used to represent the shape of the output in an affine transformation, in the format ``[N, C, H, W]`` or ``[N, C, D, H, W]``.
When the format is ``[N, C, H, W]``, it represents the batch size, number of channels, height and width. When the format is ``[N, C, D, H, W]``, it represents the batch size, number of channels, depth, height and width.
The data type must be int32.
align_corners(bool, optional): if True, aligns the centers of the 4 (4D) or 8 (5D) corner pixels of the input and output tensors, and preserves the value of the corner pixels. Default: True
name(str|None, optional): The default value is None. Normally there is no need for user to set this property. For more information, please refer to :ref:`api_guide_Name`.
Returns:
Tensor, A Tensor with shape [batch_size, H, W, 2] or [batch, D, H, W, 3] while ('D')'H', 'W' are the (depth)height, width of feature map in affine transformation. The data type is the same as `theta`.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn.functional as F
>>> # theta.shape = [1, 2, 3]
>>> theta = paddle.to_tensor(
... [
... [
... [-0.7, -0.4, 0.3],
... [0.6, 0.5, 1.5],
... ]
... ],
... dtype="float32",
... )
>>> y_t = F.affine_grid(
... theta,
... [1, 2, 3, 3],
... align_corners=False,
... )
>>> print(y_t)
Tensor(shape=[1, 3, 3, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[ 1.03333330, 0.76666665],
[ 0.56666672, 1.16666663],
[ 0.10000002, 1.56666672]],
[[ 0.76666665, 1.09999990],
[ 0.30000001, 1.50000000],
[-0.16666666, 1.90000010]],
[[ 0.50000000, 1.43333328],
[ 0.03333333, 1.83333337],
[-0.43333334, 2.23333335]]]])
"""
if not isinstance(theta, (Variable, paddle.pir.Value)):
raise TypeError("The theta should be a Tensor.")
cudnn_version = get_cudnn_version()
if cudnn_version is not None and cudnn_version >= 6000 and align_corners:
use_cudnn = True
else:
use_cudnn = False
if theta.shape[1] == 3:
use_cudnn = False
if is_compiled_with_rocm():
use_cudnn = (
False # ROCM platform do not have MIOPEN kernel for affine_grid
)
if paddle.get_flags(["FLAGS_use_accuracy_compatible_kernel"]).get(
"FLAGS_use_accuracy_compatible_kernel", False
):
use_cudnn = False
if in_dynamic_mode():
_out_shape = (
out_shape.tolist() if isinstance(out_shape, Variable) else out_shape
)
if isinstance(_out_shape, paddle.Tensor) and _out_shape.size == 0:
raise ValueError("The out_shape cannot be empty.")
theta = theta._use_gpudnn(use_cudnn)
return _C_ops.affine_grid(theta, _out_shape, align_corners)
elif in_pir_mode():
return _C_ops.affine_grid(
theta,
out_shape,
align_corners,
)
else:
helper = LayerHelper('affine_grid', **locals())
check_variable_and_dtype(
theta, 'theta', ['float32', 'float64'], 'affine_grid'
)
out = helper.create_variable_for_type_inference(dtype=theta.dtype)
inputs = {'Theta': theta}
attrs = {"align_corners": align_corners, "use_cudnn": use_cudnn}
if isinstance(out_shape, Variable):
inputs['OutputShape'] = out_shape
check_variable_and_dtype(
out_shape, 'out_shape', ['int32'], 'affine_grid'
)
else:
attrs['output_shape'] = out_shape
helper.append_op(
type='affine_grid',
inputs=inputs,
outputs={'Output': out},
attrs=None if len(attrs) == 0 else attrs,
)
return out
@param_one_alias(["x", "input"])
def pixel_unshuffle(
x: Tensor,
downscale_factor: int,
data_format: DataLayout2D = 'NCHW',
name: str | None = None,
) -> Tensor:
"""
This API implements pixel unshuffle operation.
See more details in :ref:`PixelUnShuffle <api_paddle_nn_PixelUnshuffle>` .
Parameters:
x (Tensor): 4-D tensor, the data type should be float32 or float64.
Alias: ``input``.
downscale_factor (int): Factor to decrease spatial resolution.
data_format (str, optional): The data format of the input and output data. An optional string of ``'NCHW'`` or ``'NHWC'``. When it is ``'NCHW'``, the data is stored in the order of [batch_size, input_channels, input_height, input_width]. Default: ``'NCHW'``.
name (str|None, optional): Name for the operation (optional, default is None). Normally there is no need for user to set this property. For more information, please refer to :ref:`api_guide_Name`.
Returns:
Out (Tensor): Reshaped tensor according to the new dimension.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn.functional as F
>>> x = paddle.randn([2, 1, 12, 12])
>>> out = F.pixel_unshuffle(x, 3)
>>> print(out.shape)
paddle.Size([2, 9, 4, 4])
"""
if len(x.shape) != 4:
raise ValueError(
f"Input x should be 4D tensor, but received x with the shape of {x.shape}"
)
if not isinstance(downscale_factor, int):
raise TypeError("Downscale factor must be int type")
if downscale_factor <= 0:
raise ValueError("Downscale factor must be positive")
if data_format not in ["NCHW", "NHWC"]:
raise ValueError(
"Attr(data_format) should be 'NCHW' or 'NHWC'."
f"But receive Attr(data_format): {data_format} "
)
if in_dynamic_or_pir_mode():
return _C_ops.pixel_unshuffle(x, downscale_factor, data_format)
helper = LayerHelper("pixel_unshuffle", **locals())
check_variable_and_dtype(
x, 'x', ['float16', 'float32', 'float64', 'uint16'], 'pixel_unshuffle'
)
out = helper.create_variable_for_type_inference(dtype=x.dtype)
helper.append_op(
type="pixel_unshuffle",
inputs={"X": x},
outputs={"Out": out},
attrs={
"downscale_factor": downscale_factor,
"data_format": data_format,
},
)
return out
def channel_shuffle(
x: Tensor,
groups: int,
data_format: DataLayout2D = 'NCHW',
name: str | None = None,
) -> Tensor:
"""
This API implements channel shuffle operation.
See more details in :ref:`api_paddle_nn_ChannelShuffle`.
Parameters:
x (Tensor): 4-D tensor, the data type should be float32 or float64.
groups (int): Number of groups to divide channels in.
data_format (str, optional): The data format of the input and output data. An optional string of NCHW or NHWC. The default is NCHW. When it is NCHW, the data is stored in the order of [batch_size, input_channels, input_height, input_width].
name (str|None, optional): Name for the operation (optional, default is None). Normally there is no need for user to set this property. For more information, please refer to :ref:`api_guide_Name`.
Returns:
Out (Tensor): Rearranged tensor keeping the original tensor shape.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.nn.functional as F
>>> x = paddle.arange(0, 0.6, 0.1, 'float32')
>>> x = paddle.reshape(x, [1, 6, 1, 1])
>>> print(x)
Tensor(shape=[1, 6, 1, 1], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[0. ]],
[[0.10000000]],
[[0.20000000]],
[[0.30000001]],
[[0.40000001]],
[[0.50000000]]]])
>>> y = F.channel_shuffle(x, 3)
>>> print(y)
Tensor(shape=[1, 6, 1, 1], dtype=float32, place=Place(cpu), stop_gradient=True,
[[[[0. ]],
[[0.20000000]],
[[0.40000001]],
[[0.10000000]],
[[0.30000001]],
[[0.50000000]]]])
"""
if len(x.shape) != 4:
raise ValueError(
f"Input x should be 4D tensor, but received x with the shape of {x.shape}"
)
if not isinstance(groups, int):
raise TypeError("groups must be int type")
if groups <= 0:
raise ValueError("groups must be positive")
if data_format not in ["NCHW", "NHWC"]:
raise ValueError(
"Attr(data_format) should be 'NCHW' or 'NHWC'."
f"But receive Attr(data_format): {data_format} "
)
if in_dynamic_or_pir_mode():
return _C_ops.channel_shuffle(x, groups, data_format)
helper = LayerHelper("channel_shuffle", **locals())
check_variable_and_dtype(x, 'x', ['float32', 'float64'], 'channel_shuffle')
out = helper.create_variable_for_type_inference(dtype=x.dtype)
helper.append_op(
type="channel_shuffle",
inputs={"X": x},
outputs={"Out": out},
attrs={"groups": groups, "data_format": data_format},
)
return out