pyTorch¶

class transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs)¶

将线性变换应用于传入数据 \(y = xA^T + b\)

在 NVIDIA GPU 上，它是 torch.nn.Linear 的直接替代品。

参数

in_features (int) – 每个输入样本的大小。
out_features (int) – 每个输出样本的大小。
bias (bool, default = True) – 如果设置为 False，该层将不学习加性偏差。
init_method (Callable, default = None) – 用于以下列方式初始化权重：init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
parameters_split (Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None) – 如果提供了字符串元组或字符串到整数的字典，则模块的权重和偏差参数将作为 N 个单独的 torch.nn.parameter.Parameter`s 每个参数公开，沿第一维度分割，其中 `N 是参数的长度，包含的字符串是分割参数的名称。在元组的情况下，每个参数都具有相同的形状。在字典的情况下，这些值给出了每个投影的 out_features。
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

并行参数

sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
tp_size (int, default = 1) – 当在初始化期间未形成 TP 组时，用作 TP (张量并行) 世界大小。在这种情况下，用户必须在正向传递之前调用已初始化模块上的 set_tensor_parallel_group(tp_group) 方法，以提供张量和序列并行集合所需的张量并行组。
parallel_mode ({None, ‘Column’, ‘Row’}, default = None) – 用于决定此 Linear 层是列并行 Linear 还是行并行 Linear，如此处所述。当设置为 None 时，不执行任何通信。

优化参数

fuse_wgrad_accumulation (bool, default = ‘False’) – 如果设置为 True，则启用权重梯度的创建和累积的融合。启用后，假定权重具有额外的 main_grad 属性（代替常规的 grad 使用），该属性是预先分配的缓冲区，其大小正确，可以累积梯度。
return_bias (bool, default = False) – 当设置为 True 时，此模块本身将不应用加性偏差，而是在正向传递期间返回偏差值以及线性变换的输出 \(y = xA^T\)。当偏差加法可以融合到后续操作时，这很有用。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。

forward(inp: torch.Tensor, is_first_microbatch: Optional[bool] = None)¶

将线性变换应用于输入。

参数

inp (torch.Tensor) – 输入张量。
is_first_microbatch ({True, False, None}, default = None) –
在使用梯度累积或流水线并行进行训练期间，会将数据的微型批次进一步拆分为微批次。在同一微型批次的不同微批次之间，模型权重不会更新。设置此参数表示当前微批次是否是微型批次中的第一个。设置后，此参数将启用其他优化
- 在 FP8 训练期间，它允许缓存权重的 FP8 版本
- 它还允许跳过第一个微批次的梯度累积（因为它是产生的第一个梯度）

set_tensor_parallel_group(tp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None])¶

在执行正向传递之前，为给定模块设置张量并行组。

参数: tp_group (ProcessGroup, default = None) – 张量并行进程组。

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

对输入的小批量数据应用层归一化，如论文 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) – 添加到层归一化分母的值，用于数值稳定性。
sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。
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\]
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

类 transformer_engine.pytorch.RMSNorm(hidden_size, eps=1e-5, **kwargs)¶

对输入的小批量数据应用均方根层归一化，如论文 Root Mean Square Layer Normalization 中所述。

\[y = \frac{x}{RMS_\varepsilon(x)} * \gamma\]

其中

\[RMS_\varepsilon(x) = \sqrt{\frac{1}{n}\sum_{i=0}^nx_i^2 + \varepsilon}\]

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

参数

hidden_size (int) – 每个输入样本的大小。
eps (float, default = 1e-5) – 添加到层归一化分母的值，用于数值稳定性。
sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。
zero_centered_gamma (bool, default = 'False') –
如果设置为 ‘True’，则 RMSNorm 中的 gamma 参数初始化为 0，并且 RMSNorm 公式更改为

\[y = \frac{x}{RMS(x) + \varepsilon} * (1 + \gamma)\]
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

类 transformer_engine.pytorch.LayerNormLinear(in_features, out_features, eps=1e-5, bias=True, **kwargs)¶

将层归一化应用于传入数据，然后进行线性变换。

参数

in_features (int) – 每个输入样本的大小。
out_features (int) – 每个输出样本的大小。
eps (float, default = 1e-5) – 添加到层归一化分母的值，用于数值稳定性。
bias (bool, default = True) – 如果设置为 False，该层将不学习加性偏差。
normalization ({ 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm') – 应用的归一化类型。
init_method (Callable, default = None) – 用于以下列方式初始化权重：init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
return_layernorm_output (bool, default = False) – 如果设置为 True，则 layernorm 的输出将与线性变换的输出一起从 forward 传递返回。示例用例：transformer 模块的残差连接是在 layernorm 之后进行的。
parameters_split (Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None) – 如果提供了字符串元组或字符串到整数的字典，则模块的权重和偏差参数将作为 N 个单独的 torch.nn.parameter.Parameter`s 每个参数公开，沿第一维度分割，其中 `N 是参数的长度，包含的字符串是分割参数的名称。在元组的情况下，每个参数都具有相同的形状。在字典的情况下，这些值给出了每个投影的 out_features。
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\]
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

并行参数

sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
tp_size (int, default = 1) – 当在初始化期间未形成 TP 组时，用作 TP (张量并行) 世界大小。在这种情况下，用户必须在正向传递之前调用已初始化模块上的 set_tensor_parallel_group(tp_group) 方法，以提供张量和序列并行集合所需的张量并行组。
parallel_mode ({None, ‘Column’, ‘Row’}, default = None) – 用于决定此 Linear 层是列并行 Linear 还是行并行 Linear，如此处所述。当设置为 None 时，不执行任何通信。

优化参数

fuse_wgrad_accumulation (bool, default = ‘False’) – 如果设置为 True，则启用权重梯度的创建和累积的融合。启用后，假定权重具有额外的 main_grad 属性（代替常规的 grad 使用），该属性是预先分配的缓冲区，其大小正确，可以累积梯度。
return_bias (bool, default = False) – 当设置为 True 时，此模块本身将不应用加性偏差，而是在正向传递期间返回偏差值以及线性变换的输出 \(y = xA^T\)。当偏差加法可以融合到后续操作时，这很有用。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。

forward(inp: torch.Tensor, is_first_microbatch: Optional[bool] = None)¶

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

参数

inp (torch.Tensor) – 输入张量。
is_first_microbatch ({True, False, None}, default = None) –
在使用梯度累积或流水线并行进行训练期间，会将数据的微型批次进一步拆分为微批次。在同一微型批次的不同微批次之间，模型权重不会更新。设置此参数表示当前微批次是否是微型批次中的第一个。设置后，此参数将启用其他优化
- 在 FP8 训练期间，它允许缓存权重的 FP8 版本
- 它还允许跳过第一个微批次的梯度累积（因为它是产生的第一个梯度）

set_tensor_parallel_group(tp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None])¶

在执行正向传递之前，为给定模块设置张量并行组。

参数: tp_group (ProcessGroup, default = None) – 张量并行进程组。

类 transformer_engine.pytorch.LayerNormMLP(hidden_size, ffn_hidden_size, eps=1e-5, bias=True, **kwargs)¶

将层归一化应用于输入，然后应用 MLP 模块，该模块由 2 个连续的线性变换组成，由 GeLU 激活分隔。

参数

hidden_size (int) – 每个输入样本的大小。
ffn_hidden_size (int) – 输入样本投影到的中间大小。
eps (float, default = 1e-5) – 添加到层归一化分母的值，用于数值稳定性。
bias (bool, default = True) – 如果设置为 False，则 FC1 和 FC2 层不会学习加性偏差。
normalization ({ 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm') – 应用的归一化类型。
activation (str, default = 'gelu') – 使用的激活函数。选项：‘gelu’、‘geglu’、‘relu’、‘reglu’、‘squared_relu’、‘swiglu’。
init_method (Callable, default = None) – 用于以下方式初始化 FC1 权重：init_method(weight)。如果设置为 None，则默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
output_layer_init_method (Callable, default = None) – 用于以下方式初始化 FC2 权重：output_layer_init_method(weight)。如果设置为 None，则默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
return_layernorm_output (bool, default = False) – 如果设置为 True，则 layernorm 的输出将与线性变换的输出一起从 forward 传递返回。示例用例：transformer 模块的残差连接是在 layernorm 之后进行的。
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\]
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

并行参数

set_parallel_mode (bool, default = False) – 如果设置为 True，则 FC1 用作列并行，FC2 用作行并行，如这里所述。
sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
tp_size (int, default = 1) – 当在初始化期间未形成 TP 组时，用作 TP (张量并行) 世界大小。在这种情况下，用户必须在正向传递之前调用已初始化模块上的 set_tensor_parallel_group(tp_group) 方法，以提供张量和序列并行集合所需的张量并行组。

优化参数

fuse_wgrad_accumulation (bool, default = ‘False’) – 如果设置为 True，则启用权重梯度的创建和累积的融合。启用后，假定权重具有额外的 main_grad 属性（代替常规的 grad 使用），该属性是预先分配的缓冲区，其大小正确，可以累积梯度。
return_bias (bool, default = False) – 当设置为 True 时，此模块将不对 FC2 应用加性偏差，而是在前向传递过程中返回偏差值，以及线性变换的输出 \(y = xA^T\)。当偏差加法可以融合到后续操作时，这非常有用。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。
seq_length (int) – 输入样本的序列长度。JIT Warmup 需要此参数，JIT Warmup 是一种在训练之前预热 jit 融合函数的技术，以确保正向传播和激活重新计算阶段使用相同的内核。
micro_batch_size (int) – 每个训练步骤的批量大小。JIT Warmup 需要此参数，JIT Warmup 是一种在训练之前预热 jit 融合函数的技术，以确保正向传播和激活重新计算阶段使用相同的内核。

forward(inp: torch.Tensor, is_first_microbatch: Optional[bool] = None)¶

将层归一化应用于输入，然后应用前馈网络（MLP 块）。

参数

inp (torch.Tensor) – 输入张量。
is_first_microbatch ({True, False, None}, default = None) –
在使用梯度累积或流水线并行进行训练期间，会将数据的微型批次进一步拆分为微批次。在同一微型批次的不同微批次之间，模型权重不会更新。设置此参数表示当前微批次是否是微型批次中的第一个。设置后，此参数将启用其他优化
- 在 FP8 训练期间，它允许缓存权重的 FP8 版本
- 它还允许跳过第一个微批次的梯度累积（因为它是产生的第一个梯度）

set_tensor_parallel_group(tp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None])¶

在执行正向传递之前，为给定模块设置张量并行组。

参数: tp_group (ProcessGroup, default = None) – 张量并行进程组。

类 transformer_engine.pytorch.DotProductAttention(num_attention_heads, kv_channels, **kwargs)¶

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

注意

forward 调用中的参数 attention_mask 仅当 attn_mask_type 包括 ‘“padding”’ 或 “arbitrary” 时使用。

警告

FlashAttention 使用非确定性算法来实现最佳性能。为了以牺牲性能为代价来观察确定性行为，请使用 FlashAttention 版本 < 2.0.0 并设置环境变量 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0。为了完全禁用 `flash-attn`，请设置 NVTE_FLASH_ATTN=0。

参数

num_attention_heads (int) – transformer 层中的注意力头数。
kv_channels (int) – 键值通道的数量。
num_gqa_groups (Optional[int] = None) – transformer 层中 GQA 组的数量。分组查询注意力在这篇论文中进行了描述。这只会影响键和值，而不会影响查询。GQA-1 等效于多查询注意力 (MQA)，而 GQA-H 等效于 MHA，即 num_gqa_groups = num_attention_heads。
attention_dropout (float, default = 0.0) – 多头注意力期间 dropout 运算的 dropout 概率。
attn_mask_type (str, 默认 = causal) – 传递到 softmax 操作中的 attention mask 类型，选项包括 “no_mask”, “padding”, “causal”, “padding,causal”, “causal,padding” 和 “arbitrary”，其中 “padding,causal” 和 “causal,padding” 等价。该参数可以被 attn_mask_type 在 forward 方法中覆盖。它对于涉及编译/追踪的情况很有用，例如 ONNX 导出，而 forward 参数对于动态更改 mask 类型很有用，例如，训练和推理使用不同的 mask。对于 “no_mask”，不应用任何 attention mask。对于 “causal” 或 “padding,causal” 中的 causal mask，TransformerEngine 会计算并应用一个上三角 mask 到 softmax 输入。不需要用户输入。对于 “padding” 或 “padding,causal” 中的 padding mask，用户需要通过 [batch_size + 1] 形状的 cu_seqlens_q 和 cu_seqlens_kv 或者 [batch_size, 1, 1, max_seq_len] 形状的 attention_mask 来提供 padding tokens 的位置。对于 “arbitrary” mask，用户需要提供一个可以广播到 softmax 输入形状的 mask。
attention_type (str, 默认 = self) – attention 的类型，可以是 “self” 或 “cross”。
layer_number (int, 默认 = None) – 当前 DotProductAttention 的层数，当多个这样的模块被连接在一起时，例如在连续的 transformer blocks 中。
qkv_format (str, 默认 = sbhd) – query_layer、key_layer 和 value_layer 的维度格式，{sbhd, bshd, thd}。s 代表序列长度，b 代表 batch size，h 代表 head 的数量，d 代表 head size，t 代表一个 batch 中的序列总数，其中 t = sum(s_i), for i = 0…b-1。sbhd 和 bshd 格式用于 batch 中的序列长度相等或 padding 到相等长度的情况，而 thd 格式用于 batch 中的序列具有不同长度的情况。请注意，这些格式不反映 tensors query_layer、key_layer、value_layer 在内存中的布局。为此，请使用 _get_qkv_layout 获取布局信息。

并行参数

sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
tp_size (int, default = 1) – tensor 并行 world size。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
cp_group (ProcessGroup, 默认 = None) – context 并行 process group。
cp_global_ranks (global rank IDs 列表, 默认 = None) – 在 cp_group 中的 GPUs 的 global rank IDs。
cp_stream (CUDA stream, 默认 = None) – context 并行将 flash attention 分割成多个步骤，用于计算和通信重叠。为了解决每个分割步骤的波量化问题，我们添加了一个额外的 CUDA stream，以便我们可以重叠两个 flash attention kernels。

forward(query_layer: torch.Tensor, key_layer: torch.Tensor, value_layer: torch.Tensor, attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, qkv_format: Optional[str] = None, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_kv: Optional[torch.Tensor] = None, attn_mask_type: Optional[str] = None, checkpoint_core_attention: bool, core_attention_bias_type: str = 'no_bias', core_attention_bias: Optional[torch.Tensor] = None, fast_zero_fill: bool = True)¶

点积 Attention 层。

注意

参数 attention_mask 仅在 attn_mask_type 包含 '“padding”' 或 “arbitrary” 时使用。

注意

输入 tensors query_layer, key_layer, 和 value_layer 必须都具有形状 (sequence_length, batch_size, num_attention_heads, kv_channels)。返回形状 (sequence_length, batch_size, num_attention_heads * kv_channels) 的输出。

注意

DotProductAttention 支持三种后端：1) FlashAttention，它调用 HazyResearch/Dao-AILab 的 flash-attn PyTorch API，2) FusedAttention，它有多个基于 cuDNN Graph API 的 fused attention 实现（有关 FusedAttention 后端的更多详细信息，请参阅 FusedAttention），以及 3) UnfusedDotProductAttention，它是具有 fused scaled masked softmax 的原生 PyTorch 实现。

注意

用户可以使用环境变量 NVTE_FLASH_ATTN、NVTE_FUSED_ATTN 和 NVTE_FUSED_ATTN_BACKEND 来控制使用哪个 DotProductAttention 后端，以及如果适用，使用哪个 FusedAttention 后端。TransformerEngine 优先使用 FlashAttention 而不是 FusedAttention，并优先使用 FusedAttention 而不是 UnfusedDotProductAttention。如果正在使用 FusedAttention，用户还可以通过设置 NVTE_FUSED_ATTN_USE_FAv2_BWD=1 (默认: 0) 切换到 flash-attn 的 backward 实现，因为 flash-attn 的各个版本和 FusedAttention 之间存在性能差异。此外，NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT 可用于启用 (1) 或禁用 (0) FusedAttention 中与 workspace 相关的优化。当未设置时，TransformerEngine 会根据其内部逻辑确定代码路径。这些优化以内存换取性能，应谨慎使用。

参数

query_layer (torch.Tensor) – Query tensor。
key_layer (torch.Tensor) – Key tensor。
value_layer (torch.Tensor) – Value tensor。
attention_mask (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],) – 默认 = None。用于 mask out attention softmax 输入的布尔型 tensor(s)。对于 ‘causal’ 和 ‘no_mask’ 类型，它应该为 ‘None’。对于 ‘padding’ masks，对于 self-attention，它应该是一个形状为 [batch_size, 1, 1, seqlen_q] 的单个 tensor，对于 cross-attention，它应该是一个形状为 [batch_size, 1, 1, seqlen_q] 和 [batch_size, 1, 1, seqlen_kv] 的两个 tensors 的元组。对于 ‘arbitrary’ mask 类型，它应该是一个可以广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] 形状的 tensor。
qkv_format (str, 默认 = None) – 如果提供，则覆盖从初始化来的 qkv_format。
cu_seqlens_q (Optional[torch.Tensor], 默认 = None) – 用于 query_layer 的 batch 中序列长度的累积和，形状为 [batch_size + 1]，dtype 为 torch.int32。
cu_seqlens_kv (Optional[torch.Tensor], 默认 = None) – 用于 key_layer 和 value_layer 的 batch 中序列长度的累积和，形状为 [batch_size + 1]，dtype 为 torch.int32。
attn_mask_type ({no_mask, padding, causal, padding,causal, causal,padding,) – arbitrary}, 默认 = None。传递到 softmax 操作中的 attention mask 类型。 ‘padding,causal’ 和 ‘causal,padding’ 是等价的。
checkpoint_core_attention (bool, 默认值 = False) – 如果为真，则在前向传播过程中重新计算 attention 的激活值，以便在反向传播过程中节省内存，否则需要占用内存来存储前向激活值直到反向传播。
core_attention_bias_type (str, 默认值 = no_bias) – 偏置类型，{no_bias, pre_scale_bias, post_scale_bias, alibi}
core_attention_bias (Optional[torch.Tensor], 默认值 = None) – Q * K.T 的偏置张量，形状为 [1, num_head, max_seqlen_q, max_seqlen_kv]。对于 ‘no_bias’ 和 ‘alibi’ 偏置类型，它应该为 ‘None’。
fast_zero_fill (bool, 默认值 = True) – 是否使用快速路径将输出张量设置为 0。

set_context_parallel_group(cp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None], cp_global_ranks: List[int], cp_stream: torch.cuda.Stream)¶

在执行前向传播之前，设置给定模块的上下文并行属性。

参数

cp_group (ProcessGroup) – 上下文并行进程组。
cp_global_ranks (List[int]) – 上下文组中的全局 ranks 列表。
cp_stream (torch.cuda.Stream) – 用于上下文并行执行的 cuda stream。

class transformer_engine.pytorch.MultiheadAttention(hidden_size, num_attention_heads, **kwargs)¶

多头注意力 (MHA)，包括 Query、Key、Value 和 Output projection。

注意

forward 调用中的参数 attention_mask 仅当 attn_mask_type 包括 ‘“padding”’ 或 “arbitrary” 时使用。

参数

hidden_size (int) – 每个输入样本的大小。
num_attention_heads (int) – transformer 层中的注意力头数。
kv_channels (int, 默认值 = None) – key-value 通道的数量。如果 None，则默认为 hidden_size / num_attention_heads。
attention_dropout (float, 默认值 = 0.1) – 多头注意力期间 dropout 操作的 dropout 概率。
layernorm_epsilon (float, 默认值 = 1e-5) – 添加到 layer normalization 分母中的一个值，用于数值稳定性。
init_method (Callable, 默认值 = None) – 用于以下方式初始化 QKV 和 FC1 权重的函数：init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
output_layer_init_method (Callable, 默认值 = None) – 用于以下方式初始化 PROJ 和 FC2 权重的函数：output_layer_init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
layer_number (int, 默认值 = None) – 当前 TransformerLayer 的层数，当多个此类模块连接以形成 transformer block 时。
attn_mask_type ({'no_mask', 'padding', 'causal', 'padding_causal' 'arbitrary'},) – 默认值 = causal 传递到 softmax 操作中的 attention mask 的类型。被 forward 方法中的 attn_mask_type 覆盖。 forward 参数对于动态更改 mask 类型很有用，例如，训练和推理使用不同的 mask。 init 参数对于涉及编译/跟踪的情况很有用，例如 ONNX 导出。
num_gqa_groups (int, 默认值 = None) –
transformer layer 中的 GQA groups 的数量。分组查询注意力在这篇论文中描述。这只会影响 keys 和 values，而不会影响 querys。 GQA-1 等效于 Multi-Query Attention (MQA)，而 GQA-H 等效于 MHA，即 num_gqa_groups = num_attention_heads。
return_layernorm_output (bool, default = False) – 如果设置为 True，则 layernorm 的输出将与线性变换的输出一起从 forward 传递返回。示例用例：transformer 模块的残差连接是在 layernorm 之后进行的。
input_layernorm (bool, 默认值 = True) – 如果设置为 False，则不对输入应用 layer normalization。
attention_type ({ 'self', 'cross' }, 默认值 = 'self') – 应用的 attention 类型。
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\]
normalization ({ 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm') – 应用的归一化类型。
qkv_weight_interleaved (bool, 默认值 = True) – 如果设置为 False，则 QKV 权重被解释为 query、key 和 value 权重沿 0th 维度的连接。默认解释是每个 attention head 的各个 q、k 和 v 权重是交错的。当使用 fuse_qkv_params=False 时，此参数设置为 False。
bias (bool, 默认值 = True) – 如果设置为 False，则 transformer layer 将不会学习任何附加的偏置。
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

并行参数

set_parallel_mode (bool, 默认值 = False) – 如果设置为 True，则 QKV 和 FC1 layers 用作 Column Parallel，而 PROJ 和 FC2 用作 Row Parallel，如此处所述。
sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
tp_size (int, default = 1) – 当在初始化期间未形成 TP 组时，用作 TP (张量并行) 世界大小。在这种情况下，用户必须在正向传递之前调用已初始化模块上的 set_tensor_parallel_group(tp_group) 方法，以提供张量和序列并行集合所需的张量并行组。

优化参数

fuse_wgrad_accumulation (bool, default = ‘False’) – 如果设置为 True，则启用权重梯度的创建和累积的融合。启用后，假定权重具有额外的 main_grad 属性（代替常规的 grad 使用），该属性是预先分配的缓冲区，其大小正确，可以累积梯度。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。
return_bias (bool, default = False) – 当设置为 True 时，此模块本身将不应用加性偏差，而是在正向传递期间返回偏差值以及线性变换的输出 \(y = xA^T\)。当偏差加法可以融合到后续操作时，这很有用。
fuse_qkv_params (bool, 默认值 = ‘False’) – 如果设置为 True，则 TransformerLayer 模块为 query-key-value 公开一个融合参数。这可以实现诸如 QKV 融合之类的优化，而无需 concatentations/splits，并且还可以启用参数 fuse_wgrad_accumulation。

forward(hidden_states: torch.Tensor, attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, encoder_output: Optional[torch.Tensor] = None, attn_mask_type: Optional[str] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[InferenceParams] = None, rotary_pos_emb: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: Optional[torch.Tensor] = None, fast_zero_fill: bool = True)¶

MultiheadAttention layer 的前向传播。

注意

参数 attention_mask 仅当 attn_mask_type 包括 “padding” 或 “arbitrary” 时使用。

参数

hidden_states (torch.Tensor) – 输入张量。
attention_mask (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],) – 默认 = None。用于 mask out attention softmax 输入的布尔型 tensor(s)。对于 ‘causal’ 和 ‘no_mask’ 类型，它应该为 ‘None’。对于 ‘padding’ masks，对于 self-attention，它应该是一个形状为 [batch_size, 1, 1, seqlen_q] 的单个 tensor，对于 cross-attention，它应该是一个形状为 [batch_size, 1, 1, seqlen_q] 和 [batch_size, 1, 1, seqlen_kv] 的两个 tensors 的元组。对于 ‘arbitrary’ mask 类型，它应该是一个可以广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] 形状的 tensor。
attn_mask_type ({'no_mask', 'padding', 'causal', 'padding_causal', 'arbitrary'},) – 默认值 = None 传递到 softmax 操作中的 attention mask 的类型。
encoder_output (Optional[torch.Tensor], 默认值 = None) – 如果使用 layer_type=”decoder”，则编码器块的输出将馈送到解码器块中。
is_first_microbatch ({True, False, None}, default = None) –
在使用梯度累积或流水线并行进行训练期间，会将数据的微型批次进一步拆分为微批次。在同一微型批次的不同微批次之间，模型权重不会更新。设置此参数表示当前微批次是否是微型批次中的第一个。设置后，此参数将启用其他优化
- 在 FP8 训练期间，它允许缓存权重的 FP8 版本
- 它还允许跳过第一个微批次的梯度累积（因为它是产生的第一个梯度）
checkpoint_core_attention (bool, 默认 = False) – 如果为 True，则在反向传播期间重新计算核心注意力机制的前向激活值，以节省原本需要存储前向激活值直到反向传播的内存。
rotary_pos_emb (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], 默认 = None) – 用于查询（query）和键（key）张量的嵌入，以应用旋转位置嵌入（rotary position embedding）。默认情况下，不应用任何输入嵌入。
core_attention_bias_type (str, 默认 = no_bias) – 偏置类型，{no_bias, pre_scale_bias, ‘post_scale_bias`, alibi}
core_attention_bias (Optional[torch.Tensor], 默认值 = None) – Q * K.T 的偏置张量，形状为 [1, num_head, max_seqlen_q, max_seqlen_kv]。对于 ‘no_bias’ 和 ‘alibi’ 偏置类型，它应该为 ‘None’。
fast_zero_fill (bool, 默认 = True) – 是否在使用前将输出张量设置为 0。

set_context_parallel_group(cp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None], cp_global_ranks: List[int], cp_stream: torch.cuda.Stream)¶

在执行前向传播之前，设置给定模块的上下文并行属性。

参数

cp_group (ProcessGroup) – 上下文并行进程组。
cp_global_ranks (List[int]) – 上下文组中的全局 ranks 列表。
cp_stream (torch.cuda.Stream) – 用于上下文并行执行的 cuda stream。

set_tensor_parallel_group(tp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None])¶

在执行正向传递之前，为给定模块设置张量并行组。

参数: tp_group (ProcessGroup, default = None) – 张量并行进程组。

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

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

注意

在 forward 调用中，参数 attention_mask 仅在 self_attn_mask_type 包含 “padding” 或 “arbitrary” 时使用。

参数

hidden_size (int) – 每个输入样本的大小。
ffn_hidden_size (int) – 输入样本投影到的中间大小。
num_attention_heads (int) – transformer 层中的注意力头数。
num_gqa_groups (int, 默认值 = None) –
transformer layer 中的 GQA groups 的数量。分组查询注意力在这篇论文中描述。这只会影响 keys 和 values，而不会影响 querys。 GQA-1 等效于 Multi-Query Attention (MQA)，而 GQA-H 等效于 MHA，即 num_gqa_groups = num_attention_heads。
layernorm_epsilon (float, 默认值 = 1e-5) – 添加到 layer normalization 分母中的一个值，用于数值稳定性。
hidden_dropout (float, 默认 = 0.1) – FC2 层之后 dropout 操作的 dropout 概率。
attention_dropout (float, 默认值 = 0.1) – 多头注意力期间 dropout 操作的 dropout 概率。
init_method (Callable, 默认值 = None) – 用于以下方式初始化 QKV 和 FC1 权重的函数：init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
output_layer_init_method (Callable, 默认值 = None) – 用于以下方式初始化 PROJ 和 FC2 权重的函数：output_layer_init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。
apply_residual_connection_post_layernorm (bool, 默认 = False) – 如果设置为 True，则残差连接取自 layer norm 的输出（默认取自 layer norm 的输入）
layer_number (int, 默认值 = None) – 当前 TransformerLayer 的层数，当多个此类模块连接以形成 transformer block 时。
output_layernorm (bool, 默认 = False) – 如果设置为 True，则在输出侧应用 layer normalization，在最后的 dropout-add 之后。默认行为是在输入侧应用 layer normalization，在 QKV 转换之前。
parallel_attention_mlp (bool, 默认 = False) – 如果设置为 True，则自注意力机制和前馈网络基于相同的输入（并行）计算，而不是顺序计算。两个块都具有独立的归一化。此架构在 Falcon 模型中使用。
layer_type ({‘encoder’, ‘decoder’}, 默认 = encoder) – 如果设置为 decoder，则在 self-attn 之后添加一个额外的 cross-attn 块。这可以与 encoder 选项结合使用，用于类似 T5 Transformer 的结构。
kv_channels (int, 默认值 = None) – key-value 通道的数量。如果 None，则默认为 hidden_size / num_attention_heads。
self_attn_mask_type ({'no_mask', 'padding', 'causal', 'padding_causal', 'arbitrary'},) – 默认 = causal 传递到 softmax 运算中的 attention mask 的类型。被 forward 方法中的 self_attn_mask_type 覆盖。 forward 参数对于动态更改 mask 类型很有用，例如，训练和推理使用不同的 mask。 init 参数对于涉及编译/追踪的情况很有用，例如 ONNX 导出。
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\]
normalization ({ 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm') – 应用的归一化类型。
qkv_weight_interleaved (bool, 默认值 = True) – 如果设置为 False，则 QKV 权重被解释为 query、key 和 value 权重沿 0th 维度的连接。默认解释是每个 attention head 的各个 q、k 和 v 权重是交错的。当使用 fuse_qkv_params=False 时，此参数设置为 False。
bias (bool, 默认值 = True) – 如果设置为 False，则 transformer layer 将不会学习任何附加的偏置。
activation (str, 默认 = 'gelu') – MLP 块中使用的激活类型。选项有：'gelu'、'relu'、'reglu'、'geglu' 和 'swiglu'。
device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保所有参数在运行前向传递之前都已移动到 GPU。

并行参数

set_parallel_mode (bool, 默认值 = False) – 如果设置为 True，则 QKV 和 FC1 layers 用作 Column Parallel，而 PROJ 和 FC2 用作 Row Parallel，如此处所述。
sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
tp_size (int, default = 1) – 当在初始化期间未形成 TP 组时，用作 TP (张量并行) 世界大小。在这种情况下，用户必须在正向传递之前调用已初始化模块上的 set_tensor_parallel_group(tp_group) 方法，以提供张量和序列并行集合所需的张量并行组。

优化参数

fuse_wgrad_accumulation (bool, default = ‘False’) – 如果设置为 True，则启用权重梯度的创建和累积的融合。启用后，假定权重具有额外的 main_grad 属性（代替常规的 grad 使用），该属性是预先分配的缓冲区，其大小正确，可以累积梯度。
params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，这很有用。
seq_length (int) – 输入样本的序列长度。JIT Warmup 需要此参数，JIT Warmup 是一种在训练之前预热 jit 融合函数的技术，以确保正向传播和激活重新计算阶段使用相同的内核。
micro_batch_size (int) – 每个训练步骤的批量大小。JIT Warmup 需要此参数，JIT Warmup 是一种在训练之前预热 jit 融合函数的技术，以确保正向传播和激活重新计算阶段使用相同的内核。
drop_path_rate (float, 默认 = 0.0) – 当 > 0.0 时，在残差块的主路径中对每个样本应用随机深度（stochastic depth）。
fuse_qkv_params (bool, 默认值 = ‘False’) – 如果设置为 True，则 TransformerLayer 模块为 query-key-value 公开一个融合参数。这可以实现诸如 QKV 融合之类的优化，而无需 concatentations/splits，并且还可以启用参数 fuse_wgrad_accumulation。

forward(hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, self_attn_mask_type: Optional[str] = None, encoder_output: Optional[torch.Tensor] = None, enc_dec_attn_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[transformer_engine.pytorch.attention.InferenceParams] = None, rotary_pos_emb: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: Optional[torch.Tensor] = None, fast_zero_fill: bool = True)¶

Transformer Layer: 注意力模块和一个前馈网络 (MLP)

注意

参数 attention_mask 仅在 self_attn_mask_type 包含 “padding” 或 “arbitrary” 时使用。

参数

hidden_states (torch.Tensor) – 输入张量。
attention_mask (Optional[torch.Tensor], 默认 = None) – 用于屏蔽自注意力 softmax 输入的布尔张量。对于 ‘padding’ 掩码，它应该在 [batch_size, 1, 1, seqlen_q] 中，并且对于 ‘arbitrary’，可以广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]。对于 ‘causal’ 和 ‘no_mask’，它应该为 ‘None’。
self_attn_mask_type ({'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'},) – 默认 = causal 传递到 softmax 运算中的 attention mask 的类型。
encoder_output (Optional[torch.Tensor], 默认值 = None) – 如果使用 layer_type=”decoder”，则编码器块的输出将馈送到解码器块中。
enc_dec_attn_mask (可选[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],) – 默认值 = None。如果使用 layer_type=”decoder”，则用于屏蔽交叉注意力 softmax 输入的布尔张量。对于 ‘padding’ 掩码，它应该是 [batch_size, 1, 1, seqlen_q] 和 [batch_size, 1, 1, seqlen_kv] 的两个掩码的元组。对于 ‘arbitrary’ 掩码，它应该可以广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]。对于 ‘causal’ 和 ‘no_mask’，它应该是 ‘None’。
is_first_microbatch ({True, False, None}, default = None) –
在使用梯度累积或流水线并行进行训练期间，会将数据的微型批次进一步拆分为微批次。在同一微型批次的不同微批次之间，模型权重不会更新。设置此参数表示当前微批次是否是微型批次中的第一个。设置后，此参数将启用其他优化
- 在 FP8 训练期间，它允许缓存权重的 FP8 版本
- 它还允许跳过第一个微批次的梯度累积（因为它是产生的第一个梯度）
checkpoint_core_attention (bool, 默认 = False) – 如果为 True，则在反向传播期间重新计算核心注意力机制的前向激活值，以节省原本需要存储前向激活值直到反向传播的内存。
rotary_pos_emb (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], 默认 = None) – 用于查询（query）和键（key）张量的嵌入，以应用旋转位置嵌入（rotary position embedding）。默认情况下，不应用任何输入嵌入。
core_attention_bias_type (str, 默认值 = no_bias) – 偏置类型，{no_bias, pre_scale_bias, post_scale_bias, alibi}
core_attention_bias (Optional[torch.Tensor], default = None) – Q * K.T 的偏置张量。
fast_zero_fill (bool, 默认 = True) – 是否在使用前将输出张量设置为 0。
inference_params (InferenceParams, default = None) – 推理参数，传递给主模型以便在推理期间有效地计算和存储上下文。

set_context_parallel_group(cp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None], cp_global_ranks: List[int], cp_stream: torch.cuda.Stream)¶

在执行前向传播之前，设置给定模块的上下文并行属性。

参数

cp_group (ProcessGroup) – 上下文并行进程组。
cp_global_ranks (List[int]) – 上下文组中的全局 ranks 列表。
cp_stream (torch.cuda.Stream) – 用于上下文并行执行的 cuda stream。

set_tensor_parallel_group(tp_group: Union[transformer_engine.pytorch.constants.dist_group_type, None])¶

在执行正向传递之前，为给定模块设置张量并行组。

参数: tp_group (ProcessGroup, default = None) – 张量并行进程组。

class transformer_engine.pytorch.InferenceParams(max_batch_size, max_sequence_length)¶

推理参数，传递给主模型以便在推理期间有效地计算和存储上下文。

参数

max_batch_size (int) – 推理期间的最大批处理大小。
max_sequence_length (int) – 推理期间的最大序列长度。

class transformer_engine.pytorch.CudaRNGStatesTracker¶

对于模型并行，需要同时存在多个 RNG 状态，以便在模型并行区域内外执行操作。此类跟踪各种 RNG 状态，并提供实用方法来维护它们并在给定 RNG 设置下执行模型的部分。使用 add 方法，基于输入 seed 初始化 CUDA RNG 状态，并将其分配给 name。稍后，通过 fork RNG 状态，我们可以执行操作并返回到我们的起始 CUDA 状态。

add(name: str, seed: int)¶

添加新的 RNG 状态。

name: str: RNG 状态的字符串标识符。
seed: int: RNG 状态的 PyTorch 种子。

fork(name: str = 'model-parallel-rng')¶

Fork CUDA RNG 状态，执行操作，并以原始状态退出。

name: str: RNG 状态的字符串标识符。

get_states()¶: 获取 RNG 状态。复制字典，以便我们拥有指向状态的直接指针，而不仅仅是指向字典的指针。

reset()¶: 设置为初始状态（无跟踪器）。

set_states(states: Dict[str, torch.Tensor])¶

设置 RNG 状态。出于效率考虑，我们不检查种子的尺寸是否兼容。

states: Dict[str, torch.Tensor]: 从字符串名称到 RNG 状态的映射。

transformer_engine.pytorch.fp8_autocast(enabled: bool = True, calibrating: bool = False, fp8_recipe: Optional[transformer_engine.common.recipe.DelayedScaling] = None, fp8_group: Optional[transformer_engine.pytorch.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 = True) – 是否启用 fp8
calibrating (bool, default = False) – 即使在未启用 fp8 的情况下执行，校准模式也允许收集 FP8 张量的统计信息，例如 amax 和比例数据。这对于在使用更高精度进行训练时保存可用于推理的 FP8 检查点非常有用。
fp8_recipe (recipe.DelayedScaling, default = None) – 用于 FP8 训练的 recipe。
fp8_group (torch._C._distributed_c10d.ProcessGroup, default = None) – 在每个训练步骤结束时缩减 FP8 张量的 amax 的分布式组。

transformer_engine.pytorch.fp8_model_init(enabled: bool = True)¶

用于 FP8 参数初始化的上下文管理器。

用法示例

with fp8_model_init(enabled=True):
    model = transformer_engine.pytorch.Linear(768, 768)

参数

enabled (bool, default = True) –

启用后，在此 fp8_model_init 区域内创建的 Transformer Engine 模块将仅保存其参数的 FP8 副本，而不是同时存在更高精度和 FP8 副本的默认行为。将此选项设置为 True 可能会降低内存消耗，并且对于以下情况特别有用：

使用带有主权重的优化器进行完整模型训练，其中权重的高精度副本已存在于优化器中。
推理，仅使用参数的 FP8 副本。
类似 LoRA 的微调，其中模型的主要参数不发生变化。

此功能是实验性的。

transformer_engine.pytorch.checkpoint(function: Callable, distribute_saved_activations: bool, get_cuda_rng_tracker: Callable, tp_group: transformer_engine.pytorch.constants.dist_group_type, *args: Tuple[torch.Tensor, Ellipsis], **kwargs: Dict[str, Any])¶

通过用计算换取内存来检查模型的一部分。此函数基于 torch.utils.checkpoint.checkpoint。

警告

用户有责任确保从前向和后向传递调用 function 时行为完全相同。如果产生不同的输出（例如，由于全局状态），则检查点版本在数值上将不相等。

警告

元组 args 必须仅包含张量（或 None），以便符合 PyTorch 的 save_for_backward 方法。function 必须可调用，以便使用输入 args 和 kwargs 产生有效的输出。

参数

function (Callable) – pytorch 模块，用于使用指定的 args 和 kwargs 运行前向和后向传递。
distribute_saved_activations (bool) – 如果设置为 True，则第一个张量参数在保存以用于后向传递之前，会跨指定的张量并行组 (tp_group) 分布。
get_cuda_rng_tracker (Callable) – python 可调用对象，返回 CudaRNGStatesTracker() 的实例。
tp_group (ProcessGroup, default = None) – 张量并行进程组。
args (tuple) – torch 张量的元组，作为 function 的输入。
kwargs (dict) – 字符串键的字典，作为 function 的关键字参数。

transformer_engine.pytorch.onnx_export(enabled: bool = False)¶

用于导出到 ONNX 的上下文管理器。

with onnx_export(enabled=True):
    torch.onnx.export(model)

参数: enabled (bool, default = False) – 是否启用导出