paddle¶

class transformer_engine.paddle.Linear(in_features, out_features, **kwargs)¶

对输入数据应用线性变换 \(y = xA^T + b\)

参数

in_features (int) – 每个输入样本的大小。
out_features (int) – 每个输出样本的大小。
weight_attr (Union[paddle.ParamAttr, None], default = None) – 可选的用于权重的 paddle.ParamAttr。
bias_attr (Union[paddle.ParamAttr, None, bool], default = None) – 可选的用于偏置的 paddle.ParamAttr。
backend ({‘transformer_engine’, ‘paddle’}, default = ‘transformer_engine’) – 如果设置为 ‘paddle’，则仅执行框架路径，无 FP8 优化，优化受限。

并行化参数

tp_group (ProcessGroup, default = None) – 张量并行进程组。
parallel_mode ({None, ‘Column’, ‘Row’}, default = None) – 用于决定此线性层是列并行线性还是行并行线性，详见此处。设置为 None 时，不执行通信。

forward(*args, **kwargs)¶

对输入应用线性变换。

参数: inp (paddle.Tensor) – 输入张量。

class transformer_engine.paddle.LayerNorm(hidden_size, eps=1e-5, **kwargs)¶

对输入的小批量数据应用层归一化 (Layer Normalization)，详见论文 Layer Normalization。

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma + \beta\]

\(\gamma\) 和 \(\beta\) 是大小为 hidden_size 的可学习仿射变换参数。

参数

hidden_size (int) – 每个输入样本的大小。
eps (float, default = 1e-5) – 添加到层归一化分母中以增加数值稳定性的值。
weight_attr (Union[paddle.ParamAttr, None], default = None) – 可选的用于权重的 paddle.ParamAttr。
bias_attr (Union[paddle.ParamAttr, None, bool], default = None) – 可选的用于偏置的 paddle.ParamAttr。
zero_centered_gamma (bool, default = 'False') –
如果设置为 ‘True’，LayerNorm 中的 gamma 参数将被初始化为 0，并且 LayerNorm 公式变为

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta\]
backend ({‘transformer_engine’, ‘paddle’}, default = transformer_engine) – 用于 softmax 操作的后端。

class transformer_engine.paddle.LayerNormLinear(in_features, out_features, eps=1e-5, **kwargs)¶

对输入数据应用层归一化，然后进行线性变换。

参数

in_features (int) – 每个输入样本的大小。
out_features (int) – 每个输出样本的大小。
eps (float, default = 1e-5) – 添加到层归一化分母中以增加数值稳定性的值。
weight_attr (Union[paddle.ParamAttr, None], default = None) – 可选的用于权重的 paddle.ParamAttr。
bias_attr (Union[paddle.ParamAttr, None, bool], default = None) – 可选的用于偏置的 paddle.ParamAttr。
return_layernorm_output (bool, default = False) – 如果设置为 True，则前向传播将返回层归一化的输出以及线性变换的输出。用例示例：Transformer 模块的残差连接在层归一化后获取。
zero_centered_gamma (bool, default = 'False') –
如果设置为 ‘True’，LayerNorm 中的 gamma 参数将被初始化为 0，并且 LayerNorm 公式变为

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta\]
backend ({‘transformer_engine’, ‘paddle’}, default = ‘transformer_engine’) – 如果设置为 ‘paddle’，则仅执行框架路径，无 FP8 优化，优化受限。

并行化参数

tp_group (ProcessGroup, default = None) – 张量并行进程组。
parallel_mode ({None, ‘Column’, ‘Row’}, default = None) – 用于决定此线性层是列并行线性还是行并行线性，详见此处。设置为 None 时，不执行通信。

forward(*args, **kwargs)¶

对输入应用层归一化，然后进行线性变换。

参数: inp (paddle.Tensor) – 输入张量。

class transformer_engine.paddle.LayerNormMLP(hidden_size, ffn_hidden_size, eps=1e-5, **kwargs)¶

对输入应用层归一化，然后应用 MLP 模块，该模块由两个连续的线性变换组成，中间由 GeLU 激活函数隔开。

参数

hidden_size (int) – 每个输入样本的大小。
ffn_hidden_size (int) – 输入样本被投影到的中间大小。
eps (float, default = 1e-5) – 添加到层归一化分母中以增加数值稳定性的值。
weight_attr (Union[paddle.ParamAttr, None], default = None) – 可选的用于权重的 paddle.ParamAttr。
bias_attr (Union[paddle.ParamAttr, None, bool], default = None) – 可选的用于偏置的 paddle.ParamAttr。
activation (str, default = 'gelu') – 使用的激活函数。选项：‘gelu’、‘geglu’、‘relu’、‘reglu’、‘squared_relu’、‘swiglu’。
return_layernorm_output (bool, default = False) – 如果设置为 True，则前向传播将返回层归一化的输出以及线性变换的输出。用例示例：Transformer 模块的残差连接在层归一化后获取。
zero_centered_gamma (bool, default = 'False') –
如果设置为 ‘True’，LayerNorm 中的 gamma 参数将被初始化为 0，并且 LayerNorm 公式变为

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta\]
backend ({‘transformer_engine’, ‘paddle’}, default = ‘transformer_engine’) – 如果设置为 ‘paddle’，则仅执行框架路径，无 FP8 优化，优化受限。

并行化参数

set_parallel_mode (bool, default = False) – 如果设置为 True，FC1 用作列并行，FC2 用作行并行，详见此处。
tp_group (paddle.distributed.collective.Group, default = None) – 张量并行进程组。

forward(*args, **kwargs)¶

对输入应用层归一化，然后应用前馈网络 (MLP 块)。

参数: inp (paddle.Tensor) – 输入张量。

class transformer_engine.paddle.FusedScaleMaskSoftmax(attn_mask_type, mask_func, **kwargs)¶

用于 paddle 的带融合优化的缩放和掩码 Softmax 模块。

参数

attn_mask_type (str, default = causal) – 注意力掩码类型，可以是 ‘causal’（因果）、‘padding’（填充）或 ‘no_mask’（无掩码）。
mask_func (callable) – 应用掩码到 softmax 输入的自定义可调用对象。masked_input=mask_func(inp, mask)。
softmax_in_fp32 (bool, default = True) – 在 fp32 中执行 softmax 计算。
layernorm_epsilon (float, default = 1e-5) – 添加到层归一化分母中以增加数值稳定性的值。
backend ({‘transformer_engine’, ‘paddle’}, default = transformer_engine) – 用于操作的后端。

forward(inp: paddle.Tensor, mask: paddle.Tensor, scale: Optional[float] = None)¶: FusedScaleMaskSoftmax 前向传播

class transformer_engine.paddle.DotProductAttention(num_attention_heads, kv_channels, **kwargs)¶

允许模型联合关注来自不同表示子空间的信息，详见论文：Attention Is All You Need。

注意

当 attn_mask_type 设置为 “causal” 时，forward 调用中的参数 attention_mask 将被忽略。

参数

norm_factor (float) – 注意力分数的归一化因子。
attention_dropout (float, default = 0.1) – 多头注意力操作期间的 dropout 概率。
attn_mask_type ({‘causal’, ‘padding’, ‘no_mask’}, default = causal) – 传递给 softmax 操作的注意力掩码类型。
attention_type ({‘self’, ‘cross’}, default = self) – 注意力操作类型。
backend ({‘transformer_engine’, ‘paddle’}, default = transformer_engine) – 用于注意力操作的后端。

forward(query_layer: paddle.Tensor, key_value_layer: paddle.Tensor = None, attention_mask: Optional[paddle.Tensor] = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: Optional[paddle.Tensor] = None, set_zero: bool = True)¶

点积注意力层。

注意

当 attn_mask_type 设置为 “causal” 时，参数 attention_mask 将被忽略。

注意

对于自注意力，query_layer 是沿第二个维度堆叠的 [query, key, value] 张量，其形状必须为 (batch_size, seq_length, 3, num_attention_heads, size_per_head)，且 key_value_layer 为 None。对于交叉注意力，query_layer 是 [query] 张量，其形状必须为 (batch_size, seq_length, num_attention_heads, size_per_head)，且 key_value_layer 是 [key, value] 张量，其形状必须为 (batch_size, seq_length, 2, num_attention_heads, size_per_head)。

参数

query_layer (paddle.Tensor) – Query 张量。
key_value_layer (paddle.Tensor) – Key/Value 张量。
attention_mask (Optional[paddle.Tensor], default = None) – 未使用注意力时用于掩盖 softmax 输入的布尔张量。
core_attention_bias_type (str, default = no_bias) – 当前仅支持无偏置类型，{no_bias}。
core_attention_bias (Optional[paddle.Tensor], default = None) – Q * K.T 的偏置张量。
set_zero (bool, default = True) – 是否使用快速路径将输出张量设置为 0。

class transformer_engine.paddle.MultiHeadAttention(hidden_size, num_attention_heads, **kwargs)¶

多头注意力 (MHA)，包括 Query、Key、Value 和输出投影。

参数

hidden_size (int) – 模型的隐藏大小。
num_attention_heads (int) – 注意力头的数量。
attention_dropout (float, default = 0.1) – 多头注意力操作期间的 dropout 概率。
layernorm_epsilon (float, default = 1e-5) – 层归一化操作中使用的 epsilon。
weight_attr (Union[paddle.ParamAttr, None], default = None) – 用于权重参数的 paddle.ParamAttr 对象。
bias_attr (Union[paddle.ParamAttr, None, bool], default = None) – 用于偏置参数的 paddle.ParamAttr 对象。
attn_mask_type ({‘causal’, ‘padding’, ‘no_mask’}, default = causal) – 传递给 softmax 操作的注意力掩码类型。
params_dtype (Optional[paddle.dtype], default = None) – 权重和偏置的数据类型。
return_layernorm_output (bool, default = False) – 是否返回层归一化操作的输出。
input_layernorm (bool, default = False) – 是否对输入应用层归一化。
attention_type ({‘self’, ‘cross’}, default = self) – 注意力操作类型。
zero_centered_gamma (bool, default = False) – 是否将层归一化操作的 gamma 初始化为零。
backend ({‘transformer_engine’, ‘paddle’}, default = transformer_engine) – 用于注意力操作的后端。如果设置为 ‘paddle’，则仅执行框架路径，无 FP8 优化，优化受限。

并行化参数

set_parallel_mode (bool, default = False) – 如果设置为 True，QKV 和 FC1 层用作列并行，而 PROJ 和 FC2 用作行并行，详见此处。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
rng_state_name (str, default = local_seed) – 控制用于注意力概率 dropout 的随机数生成器状态。指定的随机数生成器应为不同的 TP rank 设置不同的种子。如果 set_parallel_mode 为 False，则此项将被忽略。指定的名称应通过 paddle.distributed.fleet.meta_parallel.get_rng_state_tracker() .add(rng_state_name, seed) 进行注册。

forward(hidden_states: paddle.Tensor, attention_mask: Optional[paddle.Tensor] = None, encoder_output: Optional[paddle.Tensor] = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: Optional[paddle.Tensor] = None, set_zero: bool = True, recompute_core_attention: bool = False)¶

多头注意力层。

参数

hidden_states (paddle.Tensor) – 输入张量。
attention_mask (Optional[paddle.Tensor], default = None) – 未使用注意力时用于掩盖 softmax 输入的布尔张量。
encoder_output (Optional[paddle.Tensor], default = None) – 编码器层的输出。
core_attention_bias_type (str, default = no_bias) – 当前仅支持无偏置类型，{no_bias}。
core_attention_bias (Optional[paddle.Tensor], default = None) – Q * K.T 的偏置张量。
set_zero (bool, default = True) – 是否使用快速路径将输出张量设置为 0。
recompute_core_attention (bool, default = False) – 如果为 True，则在反向传播期间重新计算核心注意力的前向激活，以节省存储前向激活直到反向传播所需的内存。

class transformer_engine.paddle.TransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads, **kwargs)¶

TransformerLayer 由注意力块和前馈网络 (MLP) 组成。此标准层基于论文“Attention Is All You Need”。

参数

hidden_size (int) – 每个输入样本的大小。
ffn_hidden_size (int) – 输入样本被投影到的中间大小。
num_attention_heads (int) – Transformer 层中的注意力头数量。
layernorm_epsilon (float, default = 1e-5) – 添加到层归一化分母中以增加数值稳定性的值。
hidden_dropout (float, default = 0.1) – FC2 层之后的 dropout 操作的 dropout 概率。
attention_dropout (float, default = 0.1) – 多头注意力操作期间的 dropout 概率。
weight_attr (Union[paddle.ParamAttr, None], default = None) – 可选的用于权重的 paddle.ParamAttr。
bias_attr (Union[paddle.ParamAttr, None, bool], default = None) – 可选的用于偏置的 paddle.ParamAttr。
self_attn_mask_type ({‘causal’, ‘padding’}, default = causal) – 传递给 softmax 操作的自注意力掩码类型。
apply_residual_connection_post_layernorm (bool, default = False) – 如果设置为 True，残差连接将从层归一化输出中获取（默认从层归一化输入中获取）。
output_layernorm (bool, default = False) – 如果设置为 True，层归一化将在输出侧，即最终 dropout-add 之后应用。默认行为是在输入侧，即 QKV 变换之前应用层归一化。
layer_type ({‘encoder’, ‘decoder’}, default = encoder) – 如果设置为 decoder（解码器），则在自注意力之后添加一个额外的交叉注意力块。这可以与 encoder（编码器）选项结合使用，用于像 T5 Transformer 这样的结构。
zero_centered_gamma (bool, default = 'False') –
如果设置为 ‘True’，LayerNorm 中的 gamma 参数将被初始化为 0，并且 LayerNorm 公式变为

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta\]
activation (str, default = 'gelu') – MLP 块中使用的激活函数类型。选项有：‘gelu’、‘relu’、‘reglu’、‘geglu’ 和 ‘swiglu’。
params_dtype (paddle.dtype, default = paddle.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型使用较低精度训练且原始 FP32 参数无法适应 GPU 内存时非常有用。
backend ({‘transformer_engine’, ‘paddle’}, default = ‘transformer_engine’) – 如果设置为 ‘paddle’，则仅执行框架路径，无 FP8 优化，优化受限。

并行化参数

set_parallel_mode (bool, default = False) – 如果设置为 True，QKV 和 FC1 层用作列并行，而 PROJ 和 FC2 用作行并行，详见此处。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
attention_dropout_rng_state_name (str, default = local_seed) – 控制用于注意力概率 dropout 的随机数生成器状态。指定的随机数生成器应为不同的 TP rank 设置不同的种子。如果 set_parallel_mode 为 False，则此项将被忽略。
hidden_dropout_rng_state_name (str, default = global_seed) – 控制用于隐藏状态 dropout 的随机数生成器状态。指定的随机数生成器应为不同的 TP rank 设置相同的种子。如果 set_parallel_mode 为 False，则此项将被忽略。指定的名称应通过 paddle.distributed.fleet.meta_parallel.get_rng_state_tracker() .add(rng_state_name, seed) 进行注册。

forward(hidden_states: paddle.Tensor, attention_mask: Optional[paddle.Tensor] = None, encoder_output: Optional[paddle.Tensor] = None, enc_dec_attn_mask: Optional[paddle.Tensor] = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: Optional[paddle.Tensor] = None, set_zero: bool = True, recompute_core_attention: bool = False)¶

Transformer 层：注意力块和前馈网络 (MLP)

注意

当 self_attn_mask_type 设置为 “causal” 时，参数 attention_mask 将被忽略。

参数

hidden_states (paddle.Tensor) – 输入张量。
attention_mask (Optional[paddle.Tensor], default = None) – 用于掩盖自注意力 softmax 输入的布尔张量。
encoder_output (Optional[paddle.Tensor], default = None) – 如果使用 layer_type=”decoder”，则为编码器块的输出，将馈送到解码器块中。
enc_dec_attn_mask (Optional[paddle.Tensor], default = None) – 如果使用 layer_type=”decoder”，则为用于掩盖交叉注意力 softmax 输入的布尔张量。
core_attention_bias_type (str, default = no_bias) –
core_attention_bias (Optional[paddle.Tensor], default = None) – Q * K.T 的偏置张量。
set_zero (bool, default = True) – 是否在使用前将输出张量设置为 0。
recompute_core_attention (bool, default = False) – 如果为 True，则在反向传播期间重新计算核心注意力的前向激活，以节省存储前向激活直到反向传播所需的内存。

transformer_engine.paddle.fp8_autocast(enabled: bool = False, calibrating: bool = False, fp8_recipe: Optional[transformer_engine.common.recipe.DelayedScaling] = None, fp8_group: Optional[transformer_engine.paddle.constants.dist_group_type] = None)¶

FP8 使用的上下文管理器。

with fp8_autocast(enabled=True):
    out = model(inp)

注意

Transformer Engine 的 Linear 层对 FP8 的支持目前仅限于形状的两个维度都能被 16 整除的张量。对于整个 Transformer 网络的输入而言，这通常要求序列长度填充为 16 的倍数。

注意

当 fp8_recipe.reduce_amax==True 时，单个 fp8_autocast 区域内不得多次调用任何模块。这是不受支持的行为，因为 amax 缩减在 fp8_autocast 上下文退出时处理。在 fp8_autocast 区域内多次调用同一模块会在缩减发生之前覆盖 amax 张量。

参数

enabled (bool, default = False) – 是否启用 fp8
calibrating (bool, default = False) – 校准模式允许收集 fp8 张量的统计信息，例如 amax 和 scale 数据，即使在未启用 fp8 的情况下执行。这对于在使用更高精度训练时保存即用型推理 fp8 检查点非常有用。
fp8_recipe (recipe.DelayedScaling, default = None) – 用于 FP8 训练的 recipe。
fp8_group (paddle.distributed.collective.Group, default = None) – 分布式组，fp8 张量的 amax 在每个训练步骤结束时在该组上进行缩减。

transformer_engine.paddle.recompute(function, *args, **kwargs)¶

这是 paddle.distributed.fleet.utils.recompute 的包装器。它为 fp8 层提供了必要的状态信息。

参数

function (Callable) – 用于使用指定的 args 和 kwargs 运行前向和后向传播的 paddle 模块。
args (tuple) – function 输入的 torch 张量元组。
kwargs (dict) – function 关键字参数的字符串键字典。