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)。

  • get_rng_state_tracker (Callable, default = None) – 用于获取随机数生成器状态跟踪器以初始化权重。

  • rng_tracker_name (str, default = None) – 传递给 get_rng_state_tracker 以获取特定 rng 跟踪器的参数。

  • parameters_split (Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None) – 用于沿维度 0 将权重和偏置张量拆分为多个 PyTorch 参数的配置。如果提供字符串列表或元组，则它们用于创建大小相等的参数的名称。如果提供字典（最好是 OrderedDict），则键用作名称，值用作沿维度 0 的拆分大小。生成的参数的名称将以 _weight 或 _bias 结尾，因此将从任何提供的名称中删除尾随下划线。

  • 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 层是列并行线性层还是行并行线性层，如此处所述。当设置为 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: bool | None = None, fp8_output: bool | None = False, fp8_grad: bool | None = False) → torch.Tensor | Tuple[torch.Tensor, Ellipsis]

将线性变换应用于输入。

参数:

  • inp (torch.Tensor) – 输入张量。

  • is_first_microbatch ({True, False, None}, default = None) –

    在使用梯度累积或流水线并行进行训练期间，数据的小批量会进一步拆分为微批量。在同一小批量的微批量之间，模型权重不会更新。设置此参数指示当前微批量是否为小批量中的第一个。设置后，此参数启用其他优化

    • 在 FP8 训练期间，它允许缓存权重的 FP8 版本

    • 它还允许跳过第一个微批量期间的梯度累积（因为它是生成的第一个梯度）

set_tensor_parallel_group(tp_group: transformer_engine.pytorch.constants.dist_group_type | None) → None

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

参数:

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

class transformer_engine.pytorch.GroupedLinear(in_features, out_features, bias=True, **kwargs)
以分组方式对输入数据列表应用线性变换

\(y_i = x_iA_i^T + b_i\)。

参数:

  • num_gemms (int) – 要同时执行的 GEMM 的数量。

  • 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)。

  • get_rng_state_tracker (Callable, default = None) – 用于获取随机数生成器状态跟踪器以初始化权重。

  • rng_tracker_name (str, default = None) – 传递给 get_rng_state_tracker 以获取特定 rng 跟踪器的参数。

  • device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保在运行前向传播之前将所有参数移动到 GPU。

优化参数:

  • 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, m_splits: List[int], is_first_microbatch: bool | None = None) → torch.Tensor | Tuple[torch.Tensor, Ellipsis]

将线性变换应用于输入。

参数:

  • inp (torch.Tensor) – 输入张量。

  • m_splits (List[int]) – 表示输入张量拆分的整数列表。

  • is_first_microbatch ({True, False, None}, default = None) –

    在使用梯度累积或流水线并行进行训练期间，数据的小批量会进一步拆分为微批量。在同一小批量的微批量之间，模型权重不会更新。设置此参数指示当前微批量是否为小批量中的第一个。设置后，此参数启用其他优化

    • 在 FP8 训练期间，它允许缓存权重的 FP8 版本

    • 它还允许跳过第一个微批量期间的梯度累积（因为它是生成的第一个梯度）

set_tensor_parallel_group(tp_group: transformer_engine.pytorch.constants.dist_group_type | None) → None

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

参数:

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

class 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\) 是可学习的仿射变换参数，它们与输入张量的最内层维度匹配。

参数:

  • normalized_shape (int or iterable of int) – 输入张量的内部维度

  • eps (float, default = 1e-5) – 添加到层归一化分母以提高数值稳定性的值

  • device (torch.device, default = default CUDA device) – 张量设备

  • dtype (torch.dtype, default = default dtype) – 张量数据类型

  • zero_centered_gamma (bool, default = 'False') –

    如果为 True，则 \(\gamma\) 参数初始化为零，并且计算更改为

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

  • sm_margin (int or dict, default = 0) – 启动 CUDA 内核时要排除的 SM 数量。这有助于与其他内核（例如通信内核）重叠。为了更精细的控制，请提供一个字典，其中包含每个计算阶段（“forward”、“backward”、“inference”）的 SM 边距。

  • 旧版

  • ------

  • sequence_parallel (bool) – 在参数中设置名为 sequence_parallel 的 bool 属性。这是 Megatron-LM 集成的自定义逻辑。

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

均方根层归一化

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

\[y = \frac{x}{\text{RMS}_\varepsilon(x)} * \gamma\]

其中

\[\text{RMS}_\varepsilon(x) = \sqrt{\frac{1}{n}\sum_{i=0}^n x_i^2 + \varepsilon}\]

\(\gamma\) 是可学习的仿射变换参数，它与输入张量的最内层维度匹配。

参数:

  • normalized_shape (int or iterable of int) – 输入张量的内部维度

  • eps (float, default = 1e-5) – 添加到分母以提高数值稳定性的值

  • device (torch.device, default = default CUDA device) – 张量设备

  • dtype (torch.dtype, default = default dtype) – 张量数据类型

  • zero_centered_gamma (bool, default = 'False') –

    如果为 True，则 \(\gamma\) 参数初始化为零，并且计算更改为

    \[y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma)\]

  • sm_margin (int, default = 0) – 启动 CUDA 内核时要排除的 SM 数量。这有助于与其他内核（例如通信内核）重叠。为了更精细的控制，请提供一个字典，其中包含每个计算阶段（“forward”、“backward”、“inference”）的 SM 边距。

  • 旧版

  • ------

  • sequence_parallel (bool) – 在参数中设置名为 sequence_parallel 的 bool 属性。这是 Megatron-LM 集成的自定义逻辑。

class 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 的输出与线性变换的输出一起从前向传播返回。示例用例：transformer 模块的残差连接在 layernorm 后获取。

  • return_layernorm_output_gathered (bool, default = False) – 如果设置为 True，则在 all gather 操作后返回 layernorm 的输出。如果 return_layernorm_output 为 False，则忽略。示例用例：对于序列并行，transformer 模块（例如 LoRA）的残差连接的输入将需要收集。返回收集的 layernorm 输出将防止冗余收集。

  • parameters_split (Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None) – 用于沿维度 0 将权重和偏置张量拆分为多个 PyTorch 参数的配置。如果提供字符串列表或元组，则它们用于创建大小相等的参数的名称。如果提供字典（最好是 OrderedDict），则键用作名称，值用作沿维度 0 的拆分大小。生成的参数的名称将以 _weight 或 _bias 结尾，因此将从任何提供的名称中删除尾随下划线。

  • 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 层是列并行线性层还是行并行线性层，如此处所述。当设置为 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: bool | None = None, fp8_output: bool | None = False) → torch.Tensor | Tuple[torch.Tensor, Ellipsis]

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

参数:

  • inp (torch.Tensor) – 输入张量。

  • is_first_microbatch ({True, False, None}, default = None) –

    在使用梯度累积或流水线并行进行训练期间，数据的小批量会进一步拆分为微批量。在同一小批量的微批量之间，模型权重不会更新。设置此参数指示当前微批量是否为小批量中的第一个。设置后，此参数启用其他优化

    • 在 FP8 训练期间，它允许缓存权重的 FP8 版本

    • 它还允许跳过第一个微批量期间的梯度累积（因为它是生成的第一个梯度）

set_tensor_parallel_group(tp_group: transformer_engine.pytorch.constants.dist_group_type | None) → None

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

参数:

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

class 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’, ‘qgelu’, ‘srelu’。

  • 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 的输出与线性变换的输出一起从前向传播返回。示例用例：transformer 模块的残差连接在 layernorm 后获取。

  • return_layernorm_output_gathered (bool, default = False) – 如果设置为 True，则在 all gather 操作后返回 layernorm 的输出。如果 return_layernorm_output 为 False，则忽略。示例用例：对于序列并行，transformer 模块（例如 LoRA）的残差连接的输入将需要收集。返回收集的 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: bool | None = None) → torch.Tensor | Tuple[torch.Tensor, Ellipsis]

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

参数:

  • inp (torch.Tensor) – 输入张量。

  • is_first_microbatch ({True, False, None}, default = None) –

    在使用梯度累积或流水线并行进行训练期间，数据的小批量会进一步拆分为微批量。在同一小批量的微批量之间，模型权重不会更新。设置此参数指示当前微批量是否为小批量中的第一个。设置后，此参数启用其他优化

    • 在 FP8 训练期间，它允许缓存权重的 FP8 版本

    • 它还允许跳过第一个微批量期间的梯度累积（因为它是生成的第一个梯度）

set_tensor_parallel_group(tp_group: transformer_engine.pytorch.constants.dist_group_type | None) → None

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

参数:

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

class transformer_engine.pytorch.DotProductAttention(num_attention_heads, kv_channels, **kwargs)

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

注意

只有当 attn_mask_type 包含 ‘“padding”’ 或 “arbitrary” 时，才会使用 forward 调用中的参数 attention_mask。

警告

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

注意

Transformer Engine 在检查点设置时，将 FP8 元数据存储在 ._extra_state 键下。随着 FP8 注意力支持从一个后端扩展到多个后端，该键的位置也发生了变化（请参阅 FP8 检查点兼容性）。

参数:

  • num_attention_heads (int) – Transformer 层中注意力头的数量。

  • kv_channels (Union[int, Tuple[int, int]]) – 键和值张量中的头大小。如果相同，kv_channels 可以是整数；如果不同，kv_channels 应该是由两个整数组成的元组。

  • 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, default = causal) – 传递到 softmax 操作的注意力掩码类型，选项为 “no_mask”, “padding”, “causal”, “padding,causal”, “causal,padding”, “padding_causal”, “causal_bottom_right”, “padding_causal_bottom_right” 和 “arbitrary”，其中 “padding,causal”, “causal,padding” 和 “padding_causal” 是等效的。此参数可以被 forward 方法中的 attn_mask_type 覆盖。它对于涉及编译/跟踪的情况很有用，例如 ONNX 导出，而 forward 参数对于动态更改掩码类型很有用，例如，训练和推理使用不同的掩码。1. 对于 “no_mask”，不应用注意力掩码。2. 对于 “causal”, “causal_bottom_right”，或 “padding_causal” 和 “padding_causal_bottom_right” 中的因果掩码，Transformer Engine 计算并对 softmax 输入应用上三角掩码。无需用户输入。不带 “bottom_right” 后缀的因果掩码将对角线与 softmax 矩阵的左上角对齐。使用 “bottom_right”，因果掩码与右下角对齐，这通常用于推理/KV 缓存。3. 对于 “padding”，或 “padding_causal” 和 “padding_causal_bottom_right” 中的 padding 掩码，用户需要提供 padding token 的位置，可以通过 cu_seqlens_q 和 cu_seqlens_kv（形状均为 [batch_size + 1]），或者通过 attention_mask（自注意力使用形状为 [batch_size, 1, 1, max_seqlen_q] 的单个张量，交叉注意力使用形状分别为 [batch_size, 1, 1, max_seqlen_q] 和 [batch_size, 1, 1, max_seqlen_kv] 的两个张量组成的元组）。4. 对于 “arbitrary”，用户需要提供一个可以广播到 softmax 输入形状 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] 的掩码。

  • window_size (Optional[Tuple[int, int]], default = None) – 本地注意力的滑动窗口大小，位置 i 的查询关注范围为 [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]]（包含端点）。特殊情况 (-1, -1) 和 (-1, 0) 分别表示没有滑动窗口和因果掩码。 causal 和 causal_bottom_right 掩码都映射到 window_size = (-1, 0)，Transformer Engine 根据 attn_mask_type 区分它们。与 attn_mask_type 类似，window_size 也可以在 forward 中被 window_size 覆盖。

  • attention_type (str, default = self) – 注意力类型，可以是 “self” 和 “cross”。

  • layer_number (int, default = None) – 当前 DotProductAttention 的层号，当多个此类模块串联时，例如在连续的 Transformer 块中。

  • qkv_format (str, default = sbhd) – query_layer、key_layer 和 value_layer 的维度格式，{sbhd, bshd, thd}。s 代表序列长度，b 批量大小，h 注意力头的数量，d 头大小，t 批量中的 token 总数，其中 t = sum(s_i), for i = 0…b-1。sbhd 和 bshd 格式用于批次中的序列长度相等或 padding 到相等长度的情况，thd 格式用于批次中的序列长度不同的情况。请注意，这些格式不反映张量 query_layer、key_layer、value_layer 在内存中的布局方式。为此，请使用 get_qkv_layout 来获取布局信息。

  • softmax_scale (Optional[float], default = None) – 注意力分数的 softmax 缩放。如果为 None，则默认为 1.0/math.sqrt(kv_channels if isinstance(kv_channels, int) else kv_channels[0])。

并行参数:

  • sequence_parallel (bool, default = False) – 如果设置为 True，则使用序列并行。

  • tp_size (int, default = 1) – 张量并行世界大小。

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

  • cp_group (Union[ProcessGroup, List[ProcessGroup]], default = None) – 上下文并行进程组。ProcessGroup 用于 “p2p”、“all_gather” 和 “a2a” 的 cp_comm_type。List[ProcessGroup] 用于 “a2a+p2p” 的 cp_comm_type，其中 cp_group[0] 和 cp_group[1] 分别用于 a2a 和 p2p 通信。

  • cp_global_ranks (list of global rank IDs, default = None) – cp_group 中 GPU 的全局 rank ID 列表。

  • cp_stream (CUDA stream, default = None) – 上下文并行将 FlashAttention 分成多个步骤，以实现计算和通信的重叠。为了解决每个拆分步骤的波浪量化问题，我们添加了一个额外的 CUDA stream，以便我们可以重叠两个 Flash Attention 内核。

  • cp_comm_type (str, default = p2p) – 上下文并行的 GPU 间通信类型。可以是 “p2p”、“all_gather”、“a2a” 或 “a2a+p2p”。“p2p”: 在环形拓扑中使用 P2P 通信交换 KV 块。

    P2P 是异步的，可以与注意力计算重叠。

    “all_gather”: 在注意力计算之前，使用 all-gather 获取完整的 KV 序列。

    all-gather 不是异步的，不能重叠。

    “a2a”: 类似于 DeepSpeed Ulysses，将注意力头分散在 CP

    组中，并进行 gather 以获取完整的 QKV 序列。

    “a2a+p2p”: 分层 CP 实现。首先在每个 CP 子组内对 QKV 应用 a2a（例如，通过 NVLink），然后在子组之间使用 p2p 交换 KV（例如，通过 IBLink）。

forward(query_layer: torch.Tensor, key_layer: torch.Tensor, value_layer: torch.Tensor, attention_mask: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None, qkv_format: str | None = None, cu_seqlens_q: torch.Tensor | None = None, cu_seqlens_kv: torch.Tensor | None = None, cu_seqlens_q_padded: torch.Tensor | None = None, cu_seqlens_kv_padded: torch.Tensor | None = None, max_seqlen_q: int | None = None, max_seqlen_kv: int | None = None, attn_mask_type: str | None = None, window_size: Tuple[int, int] | None = None, checkpoint_core_attention: bool = False, core_attention_bias_type: str = 'no_bias', core_attention_bias: torch.Tensor | None = None, alibi_slopes: torch.Tensor | None = None, fast_zero_fill: bool = True, inference_params: InferenceParams | None = None) → torch.Tensor

点积注意力层。

注意

只有当 attn_mask_type 包含 ‘“padding”’ 或 “arbitrary” 时，才会使用参数 attention_mask。

注意

DotProductAttention 支持三个后端：1) FlashAttention，它调用 HazyResearch/Dao-AILab 的 flash-attn PyTorch API；2) FusedAttention，它有多个基于 cuDNN Graph API 的融合注意力实现（有关 FusedAttention 后端的更多详细信息，请参阅 FusedAttention）；以及 3) UnfusedDotProductAttention，它是具有融合缩放掩码 softmax 的原生 PyTorch 实现。

注意

用户可以使用环境变量 NVTE_FLASH_ATTN, NVTE_FUSED_ATTN, 和 NVTE_FUSED_ATTN_BACKEND 来控制使用哪个 DotProductAttention 后端，以及 FusedAttention 后端（如果适用）。Transformer Engine 优先使用 FlashAttention，然后是 FusedAttention，最后是 UnfusedDotProductAttention。如果正在使用 FusedAttention，用户还可以通过设置 NVTE_FUSED_ATTN_USE_FAv2_BWD=1（默认值：0）切换到 flash-attn 的反向传播实现，因为各种版本的 flash-attn 和 FusedAttention 之间存在性能差异。此外，可以使用 NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT 来启用 (1) 或禁用 (0) FusedAttention 中与工作区相关的优化。当未设置时，Transformer Engine 会根据其内部逻辑确定代码路径。这些优化以内存换取性能，应谨慎使用。

注意

当训练数据具有可变序列长度时，用户有两种选择。

  1. 操作数据并将所有序列 padding 到相同的长度。使用 qkv_format = {“bshd”, “sbhd”} 和 attn_mask_type = {“padding”, “padding_causal”, “padding_causal_bottom_right”}。传入 cu_seqlens_q 和 cu_seqlens_kv，或 attention_mask（它将被转换为 cu_seqlens_q 和 cu_seqlens_kv），以提供真实的序列长度信息。例如，一批 3 个序列 [a a a b b c c c c] 可以 padding 为 [a a a PAD b b PAD PAD c c c c]，自注意力的累积序列长度张量将为 cu_seqlens_q = cu_seqlens_kv = [0, 3, 5, 9]。

  2. 不对训练数据执行 padding。使用 qkv_format = “thd” 和 attn_mask_type = {“padding”, “padding_causal”, “padding_causal_bottom_right”}。传入 cu_seqlens_q 和 cu_seqlens_kv，或 attention_mask，如选项 1 中所述。例如，一批 3 个序列 [a a a b b c c c c] 可以无需任何 padding 进行处理，自注意力的序列长度张量将为 cu_seqlens_q = cu_seqlens_kv = [0, 3, 5, 9]。

    在某些用例中，序列之间会插入不同数量的标识符 token。这些 token 不参与注意力计算。在这种情况下，必须指定 cu_seqlens_q_padded 和 cu_seqlens_kv_padded，以正确识别批次中每个序列的开始和结束位置。例如，一批 3 个序列 [a a a 1 b b 2 2 c c c c 3] 的自注意力将具有 cu_seqlens_q = cu_seqlens_kv = [0, 3, 5, 9]，以及 cu_seqlens_q_padded = cu_seqlens_kv_padded = [0, 4, 8, 13]。

注意

当 qkv_format = {“bshd”, “sbhd”} 时，批次中的序列长度相等。max_seqlen_q 和 max_seqlen_kv 应与 query_layer 和 key_layer 张量的 “s” 维度相同。当未设置时，Transformer Engine 将推断它们为如此。

当 qkv_format = “thd” 时，序列具有不同的长度。max_seqlen_q 和 max_seqlen_kv 应为批次中最大查询和键/值序列长度。当未设置时，Transformer Engine 会从 cu_seqlens_q 和 cu_seqlens_kv 中推断它们。此推断会消耗少量内核和一些 CPU-GPU 同步开销，为了避免这种开销，建议用户从数据加载器中获取最大序列长度并将其传入。

  • 由于最大序列长度、批量大小和 token 数量在批次之间会发生变化，因此张量构造需要支持动态形状。FlashAttention 和 UnfusedDotProductAttention 自然地支持这样做，而 FusedAttention 需要参数是静态的，以便在性能启发式分析之前创建图。为了减少每次运行创建的图的数量，Transformer Engine 1.13+ 量化了相关参数：对于 cuDNN < 9.6，量化 {批量大小, max_seqlen_q, max_seqlen_kv}；对于 cuDNN >= 9.6，量化 {query_layer 的 “t” 维度, key_layer 的 “t” 维度}。

参数:

  • query_layer (torch.Tensor) – 查询张量。

  • key_layer (torch.Tensor) – 键张量。

  • value_layer (torch.Tensor) – 值张量。

  • attention_mask (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],) – default = None。用于屏蔽注意力 softmax 输入的布尔张量。对于因果掩码和 “no_mask”，它应为 None。对于 padding 掩码，自注意力应为形状为 [batch_size, 1, 1, seqlen_q] 的单个张量，交叉注意力应为形状分别为 [batch_size, 1, 1, seqlen_q] 和 [batch_size, 1, 1, seqlen_kv] 的两个张量组成的元组。对于 “arbitrary” 掩码，它应为可广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] 形状的张量。True 值表示相应位置被屏蔽，False 值表示允许该位置参与注意力计算。

  • qkv_format (str, default = None) – 如果提供，则覆盖从初始化设置的 qkv_format。

  • cu_seqlens_q (Optional[torch.Tensor], default = None) – 批次中 query_layer 的序列长度累积和（无偏移），形状为 [batch_size + 1]，数据类型为 torch.int32。有关更多详细信息，请参阅 注意。

  • cu_seqlens_kv (Optional[torch.Tensor], default = None) – 批次中 key_layer 和 value_layer 的序列长度累积和（无偏移），形状为 [batch_size + 1]，数据类型为 torch.int32。有关更多详细信息，请参阅 注意。

  • cu_seqlens_q_padded (Optional[torch.Tensor], default = None) – 批次中 query_layer 的序列长度累积和（带偏移），形状为 [batch_size + 1]，数据类型为 torch.int32。当批次中序列之间没有 padding 时，cu_seqlens_q_padded = cu_seqlens_q。有关更多详细信息，请参阅 注意。

  • cu_seqlens_kv_padded (Optional[torch.Tensor], default = None) – 批次中 key_layer 和 value_layer 的序列长度累积和（带偏移），形状为 [batch_size + 1]，数据类型为 torch.int32。当批次中序列之间没有 padding 时，cu_seqlens_kv_padded = cu_seqlens_kv。有关更多详细信息，请参阅 注意。

  • max_seqlen_q (Optional[int], default = None) – query_layer 中的最大序列长度。有关更多详细信息，请参阅 注意。

  • max_seqlen_kv (Optional[int], default = None) – key_layer 和 value_layer 中的最大序列长度。有关更多详细信息，请参阅 注意。

  • attn_mask_type ({'no_mask', 'padding', 'causal', 'padding,causal', 'causal,padding',) – ‘padding_causal’, ‘causal_bottom_right’, ‘padding_causal_bottom_right’, ‘arbitrary’}, default = None。传递到 softmax 操作的注意力掩码类型。‘padding,causal’、‘causal,padding’ 和 ‘padding_causal’ 是等效的。默认情况下，因果掩码与 softmax 矩阵的左上角对齐。当在掩码类型中指定 “bottom_right” 时，因果掩码与右下角对齐。

  • window_size (Optional[Tuple[int, int]], default = None) – 本地注意力的滑动窗口大小。

  • checkpoint_core_attention (bool, default = False) – 如果为 true，则在反向传播期间重新计算注意力的前向激活，以节省原本用于存储前向激活直到反向传播的内存。

  • core_attention_bias_type (str, default = no_bias) – 偏置类型，{no_bias, pre_scale_bias, post_scale_bias, alibi}

  • core_attention_bias (Optional[torch.Tensor], default = None) – Q * K.T 的偏置张量，形状为 [1, num_head, max_seqlen_q, max_seqlen_kv]。对于 ‘no_bias’ 和 ‘alibi’ 偏置类型，它应为 ‘None’。

  • alibi_slopes (Optional[torch.Tensor], default = None) – FP32 中的 ALiBi 斜率，形状为 [nheads] 或 [batch_size, nheads]。它将 (-alibi_slope * (i + seqlen_k - seqlen_q - j)) 的偏置添加到查询 i 和键 j 的注意力分数中。

  • fast_zero_fill (bool, default = True) – 是否使用快速路径将输出张量设置为 0。

  • inference_params (Optional[InferenceParams], default = None) – 通过缓存当前解码迭代的键和值来优化推理期间的执行性能。这些缓存值附加到先前迭代中计算的 K 和 V 值，从而无需为整个序列重新计算它们。使用前需要初始化 inference_params，以确保有足够的内存分配。序列长度偏移量的调整应在完整的正向传播后完成。如果使用旋转位置嵌入 (RoPE)，则必须事先准备好它们。支持 “sbhd” 和 “bshd” 布局，其中 “sbhd” 布局效率更高。

set_context_parallel_group(cp_group: transformer_engine.pytorch.constants.dist_group_type | List[transformer_engine.pytorch.constants.dist_group_type] | None, cp_global_ranks: List[int], cp_stream: torch.cuda.Stream, cp_comm_type: str = 'p2p') → None

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

参数:

  • cp_group (Union[ProcessGroup, List[ProcessGroup]]) – 上下文并行进程组。ProcessGroup 用于 “p2p”、“all_gather” 和 “a2a” 的 cp_comm_type。List[ProcessGroup] 用于 “a2a+p2p” 的 cp_comm_type，其中 cp_group[0] 和 cp_group[1] 分别用于 a2a 和 p2p 通信。

  • cp_global_ranks (List[int]) – 上下文组中的全局 ranks 列表。

  • cp_stream (torch.cuda.Stream) – 用于上下文并行执行的 cuda 流。

  • cp_comm_type (str, default = p2p) –

    上下文并行 inter-gpu 通信类型。 可以是 “p2p”、“all_gather”、“a2a” 或 “a2a+p2p”。 “p2p”：在环形拓扑中使用 P2P 通信交换 KV 块。

    P2P 是异步的，可以与注意力计算重叠。

    ”all_gather”：在 attention 之前，使用 all-gather 获取完整的 KV 序列。

    all-gather 不是异步的，不能重叠。

    ”a2a”：类似于 DeepSpeed Ulysses，在 CP 之间分散 attention heads

    组中，并进行 gather 以获取完整的 QKV 序列。

    ”a2a+p2p”：分层 CP 实现。首先在每个 CP 子组内对 QKV 应用 a2a（例如，通过 NVLink），然后在子组之间使用 p2p 交换 KV（例如，通过 IBLink）。

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

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

注意

只有当 attn_mask_type 包含 ‘“padding”’ 或 “arbitrary” 时，才会使用 forward 调用中的参数 attention_mask。

参数:

  • hidden_size (int) – 每个输入样本的大小。

  • num_attention_heads (int) – Transformer 层中注意力头的数量。

  • kv_channels (int, default = None) – key-value 通道的数量。如果为 None，则默认为 hidden_size / num_attention_heads。

  • attention_dropout (float, default = 0.1) – 在多头注意力期间 dropout 操作的 dropout 概率。

  • layernorm_epsilon (float, default = 1e-5) – 添加到层归一化分母中的值，用于数值稳定性。

  • init_method (Callable, default = None) – 用于初始化 QKV 和 FC1 权重的权重，方式如下：init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。

  • output_layer_init_method (Callable, default = None) – 用于初始化 PROJ 和 FC2 权重的权重，方式如下：output_layer_init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。

  • layer_number (int, default = None) – 当多个此类模块连接以形成 transformer 块时，当前 TransformerLayer 的层号。

  • attn_mask_type ({'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right',) – ‘padding_causal_bottom_right’,’arbitrary’}, default = causal 传递到 softmax 操作中的 attention mask 类型。被 forward 方法中的 attn_mask_type 覆盖。 forward 参数对于动态更改 mask 类型很有用，例如，训练和推理使用不同的 mask。 init 参数对于涉及编译/tracing 的情况很有用，例如 ONNX 导出。

  • window_size (Optional[Tuple[int, int]], default = None) – 本地注意力的滑动窗口大小，位置 i 的查询关注范围为 [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]]（包含端点）。特殊情况 (-1, -1) 和 (-1, 0) 分别表示没有滑动窗口和因果掩码。 causal 和 causal_bottom_right 掩码都映射到 window_size = (-1, 0)，Transformer Engine 根据 attn_mask_type 区分它们。与 attn_mask_type 类似，window_size 也可以在 forward 中被 window_size 覆盖。

  • num_gqa_groups (int, default = None) –

    transformer 层中 GQA 组的数量。分组查询注意力在 本文 中描述。这仅影响 keys 和 values，不影响 querys。GQA-1 等效于多查询注意力 (MQA)，而 GQA-H 等效于 MHA，即 num_gqa_groups = num_attention_heads。

  • return_layernorm_output (bool, default = False) – 如果设置为 True，则 layernorm 的输出与线性变换的输出一起从前向传播返回。示例用例：transformer 模块的残差连接在 layernorm 后获取。

  • input_layernorm (bool, default = False) – 如果设置为 True，则对输入应用层归一化。

  • attention_type ({ 'self', 'cross' }, default = '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, default = True) – 如果设置为 False，则 QKV 权重被解释为 query、key 和 value 权重沿 0th 维度的串联。默认解释是每个 attention head 的单个 q、k 和 v 权重是交错的。当使用 fuse_qkv_params=False 时，此参数设置为 False。

  • bias (bool, default = True) – 如果设置为 False，则 transformer 层将不学习任何附加偏置。

  • device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保在运行前向传播之前将所有参数移动到 GPU。

  • qkv_format (str, default = sbhd) – query_layer、key_layer 和 value_layer 的维度格式，{sbhd, bshd}。s 代表序列长度，b 批大小，h head 数量，d head 大小。sbhd 和 bshd 格式用于批次中的序列长度相等或填充到相等长度的情况。请注意，这些格式不反映张量 query_layer、key_layer、value_layer 在内存中的布局方式。为此，请使用 get_qkv_layout 获取布局信息。

并行参数:

  • set_parallel_mode (bool, default = False) – 如果设置为 True，则 QKV 和 FC1 层用作列并行，而 PROJ 和 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），该属性是预分配的缓冲区，其大小正确，可以累积梯度。

  • params_dtype (torch.dtype, default = torch.get_default_dtype()) – 它控制用于分配初始参数的类型。当模型以较低精度训练并且原始 FP32 参数不适合 GPU 内存时，此参数很有用。

  • return_bias (bool, default = False) – 当设置为 True 时，此模块将不应用加性偏置本身，而是在前向传播期间返回偏置值以及线性变换 \(y = xA^T\) 的输出。当可以将偏置加法融合到后续操作时，这很有用。

  • fuse_qkv_params (bool, default = ‘False’) – 如果设置为 True，则 TransformerLayer 模块为 query-key-value 公开单个融合参数。 这可以实现诸如无连接/拆分的 QKV 融合之类的优化，并且还启用参数 fuse_wgrad_accumulation。

forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None, encoder_output: torch.Tensor | None = None, attn_mask_type: str | None = None, window_size: Tuple[int, int] | None = None, is_first_microbatch: bool | None = None, checkpoint_core_attention: bool = False, inference_params: InferenceParams | None = None, rotary_pos_emb: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: torch.Tensor | None = None, alibi_slopes: torch.Tensor | None = None, cu_seqlens_q: torch.Tensor | None = None, cu_seqlens_kv: torch.Tensor | None = None, max_seqlen_q: int | None = None, max_seqlen_kv: int | None = None, fast_zero_fill: bool = True) → Tuple[torch.Tensor | None, Ellipsis]

MultiheadAttention 层的正向传播。

注意

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

参数:

  • hidden_states (torch.Tensor) – 输入张量。

  • attention_mask (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],) – default = None。用于屏蔽注意力 softmax 输入的布尔张量。对于因果掩码和 “no_mask”，它应为 None。对于 padding 掩码，自注意力应为形状为 [batch_size, 1, 1, seqlen_q] 的单个张量，交叉注意力应为形状分别为 [batch_size, 1, 1, seqlen_q] 和 [batch_size, 1, 1, seqlen_kv] 的两个张量组成的元组。对于 “arbitrary” 掩码，它应为可广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] 形状的张量。True 值表示相应位置被屏蔽，False 值表示允许该位置参与注意力计算。

  • attn_mask_type ({'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right',) – ‘padding_causal_bottom_right’,’arbitrary’}, default = None 传递到 softmax 操作中的 attention mask 类型。 默认情况下，因果 mask 与 softmax 矩阵的左上角对齐。 当 mask 类型中指定 “bottom_right” 时，因果 mask 与右下角对齐。

  • window_size (Optional[Tuple[int, int]], default = None) – 本地注意力的滑动窗口大小。

  • encoder_output (Optional[torch.Tensor], default = None) – 编码器块的输出，如果使用 layer_type=”decoder”，则馈送到解码器块。

  • is_first_microbatch ({True, False, None}, default = None) –

    在使用梯度累积或流水线并行进行训练期间，数据的小批量会进一步拆分为微批量。在同一小批量的微批量之间，模型权重不会更新。设置此参数指示当前微批量是否为小批量中的第一个。设置后，此参数启用其他优化

    • 在 FP8 训练期间，它允许缓存权重的 FP8 版本

    • 它还允许跳过第一个微批量期间的梯度累积（因为它是生成的第一个梯度）

  • checkpoint_core_attention (bool, default = False) – 如果为 true，则在反向传播期间重新计算核心注意力的前向激活，以节省内存，否则这些内存将被占用以存储前向激活直到反向传播。

  • rotary_pos_emb (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = None) – 用于 query 和 key 张量的嵌入，以应用旋转位置嵌入。 默认情况下，不应用输入嵌入。

  • core_attention_bias_type (str, default = no_bias) – 偏置类型，{no_bias, pre_scale_bias, ‘post_scale_bias`, alibi}

  • core_attention_bias (Optional[torch.Tensor], default = None) – Q * K.T 的偏置张量，形状为 [1, num_head, max_seqlen_q, max_seqlen_kv]。对于 ‘no_bias’ 和 ‘alibi’ 偏置类型，它应为 ‘None’。

  • alibi_slopes (Optional[torch.Tensor], default = None) – FP32 中的 ALiBi 斜率，形状为 [nheads] 或 [batch_size, nheads]。它将 (-alibi_slope * (i + seqlen_k - seqlen_q - j)) 的偏置添加到查询 i 和键 j 的注意力分数中。

  • cu_seqlens_q (Optional[torch.Tensor], default = None) – 批次中 query_layer 的序列长度的累积和（无偏移量），形状为 [batch_size + 1]，dtype 为 torch.int32。

  • cu_seqlens_kv (Optional[torch.Tensor], default = None) – 批次中 key_layer 和 value_layer 的序列长度的累积和（无偏移量），形状为 [batch_size + 1]，dtype 为 torch.int32。

  • max_seqlen_q (Optional[int], default = None) – query_layer 中的最大序列长度。 如果未提供，则从 cu_seqlens_q 计算。

  • max_seqlen_kv (Optional[int], default = None) – key_layer 和 value_layer 中的最大序列长度。 如果未提供，则从 cu_seqlens_kv 计算。

  • fast_zero_fill (bool, default = True) – 是否在使用前将输出张量设置为 0。

set_context_parallel_group(cp_group: transformer_engine.pytorch.constants.dist_group_type | List[transformer_engine.pytorch.constants.dist_group_type] | None, cp_global_ranks: List[int], cp_stream: torch.cuda.Stream, cp_comm_type: str = 'p2p') → None

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

参数:

  • cp_group (Union[ProcessGroup, List[ProcessGroup]]) – 上下文并行进程组。ProcessGroup 用于 “p2p”、“all_gather” 和 “a2a” 的 cp_comm_type。List[ProcessGroup] 用于 “a2a+p2p” 的 cp_comm_type，其中 cp_group[0] 和 cp_group[1] 分别用于 a2a 和 p2p 通信。

  • cp_global_ranks (List[int]) – 上下文组中的全局 ranks 列表。

  • cp_stream (torch.cuda.Stream) – 用于上下文并行执行的 cuda 流。

  • cp_comm_type (str, default = p2p) –

    上下文并行 inter-gpu 通信类型。 可以是 “p2p”、“all_gather”、“a2a”、“a2a+p2p”。 “p2p”：在环形拓扑中使用 P2P 通信交换 KV 块。

    P2P 是异步的，可以与注意力计算重叠。

    ”all_gather”：在 attention 之前，使用 all-gather 获取完整的 KV 序列。

    all-gather 不是异步的，不能重叠。

    ”a2a”：类似于 DeepSpeed Ulysses，在 CP 之间分散 attention heads

    组中，并进行 gather 以获取完整的 QKV 序列。

    ”a2a+p2p”：分层 CP 实现。首先在每个 CP 子组内对 QKV 应用 a2a（例如，通过 NVLink），然后在子组之间使用 p2p 交换 KV（例如，通过 IBLink）。

set_tensor_parallel_group(tp_group: transformer_engine.pytorch.constants.dist_group_type | None) → None

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

参数:

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

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

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

注意

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

参数:

  • hidden_size (int) – 每个输入样本的大小。

  • ffn_hidden_size (int) – 输入样本投影到的中间大小。

  • num_attention_heads (int) – Transformer 层中注意力头的数量。

  • num_gqa_groups (int, default = None) –

    transformer 层中 GQA 组的数量。分组查询注意力在 本文 中描述。这仅影响 keys 和 values，不影响 querys。GQA-1 等效于多查询注意力 (MQA)，而 GQA-H 等效于 MHA，即 num_gqa_groups = num_attention_heads。

  • layernorm_epsilon (float, default = 1e-5) – 添加到层归一化分母中的值，用于数值稳定性。

  • hidden_dropout (float, default = 0.1) – FC2 层之后 dropout 操作的 dropout 概率。

  • attention_dropout (float, default = 0.1) – 在多头注意力期间 dropout 操作的 dropout 概率。

  • init_method (Callable, default = None) – 用于初始化 QKV 和 FC1 权重的权重，方式如下：init_method(weight)。当设置为 None 时，默认为 torch.nn.init.normal_(mean=0.0, std=0.023)。

  • output_layer_init_method (Callable, default = 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, default = False) – 如果设置为 True，则残差连接取自层归一化的输出（默认为层归一化的输入）。

  • layer_number (int, default = None) – 当多个此类模块连接以形成 transformer 块时，当前 TransformerLayer 的层号。

  • output_layernorm (bool, default = False) – 如果设置为 True，则在最终 dropout-add 之后，在输出侧应用层归一化。 默认行为是在 QKV 转换之前的输入侧应用层归一化。

  • parallel_attention_mlp (bool, default = False) – 如果设置为 True，则自注意力和前馈网络基于相同的输入（并行）而不是顺序计算。 两个块都具有独立的归一化。 此架构用于 Falcon 模型。

  • layer_type ({‘encoder’, ‘decoder’}, default = encoder) – 如果设置为 decoder，则在自注意力之后添加一个额外的交叉注意力块。 这可以用于像 T5 Transformer 这样的结构，并结合 encoder 选项使用。

  • kv_channels (int, default = None) – 每个 attention head 的 query-key-value 通道数。 如果为 None，则默认为 hidden_size / num_attention_heads。

  • self_attn_mask_type ({'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right',) – ‘padding_causal_bottom_right’, ‘arbitrary’}, default = causal 传递到编码器的 softmax 操作中的 attention mask 类型。 被 forward 方法中的 self_attn_mask_type 覆盖。 forward 参数对于动态更改 mask 类型很有用，例如，训练和推理使用不同的 mask。 init 参数对于涉及编译/tracing 的情况很有用，例如 ONNX 导出。

  • window_size (Optional[Tuple[int, int]], default = None) – 编码器中本地注意力的滑动窗口大小，其中位置 i 的 query 注意 [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]]（包含边界）中的 keys。 特殊情况 (-1, -1) 和 (-1, 0) 分别表示没有滑动窗口和因果 mask。 causal 和 causal_bottom_right mask 都映射到 window_size = (-1, 0)，Transformer Engine 基于 self_attn_mask_type 或 enc_dec_attn_mask_type 区分它们。 与 self_attn_mask_type 类似，window_size 也可以被 forward 中的 window_size 覆盖。

  • enc_dec_attn_mask_type ({'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'},) – default = no_mask 传递到解码器的 softmax 操作中的 attention mask 类型。

  • enc_dec_window_size (Optional[Tuple[int, int]], default = None) – 解码器中本地注意力的滑动窗口大小。

  • 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, default = True) – 如果设置为 False，则 QKV 权重被解释为 query、key 和 value 权重沿 0th 维度的串联。默认解释是每个 attention head 的单个 q、k 和 v 权重是交错的。当使用 fuse_qkv_params=False 时，此参数设置为 False。

  • bias (bool, default = True) – 如果设置为 False，则 transformer 层将不学习任何附加偏置。

  • activation (str, default = 'gelu') – MLP 块中使用的激活类型。选项有：‘gelu’、‘relu’、‘reglu’、‘geglu’、‘swiglu’、‘qgelu’ 和 ‘srelu’。

  • device (Union[torch.device, str], default = "cuda") – 模型参数将分配到的设备。用户有责任确保在运行前向传播之前将所有参数移动到 GPU。

  • attn_input_format ({'sbhd', 'bshd'}, default = 'sbhd') – 这控制中间隐藏状态的维度是 ‘batch first’ (‘bshd’) 还是 ‘sequence first’ (‘sbhd’)。 s 代表序列长度，b 批大小，h head 数量，d head 大小。 请注意，这些格式与 MultiHeadAttention 和 DotProductAttention 模块中的 qkv_format 非常相关。

并行参数:

  • set_parallel_mode (bool, default = False) – 如果设置为 True，则 QKV 和 FC1 层用作列并行，而 PROJ 和 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），该属性是预分配的缓冲区，其大小正确，可以累积梯度。

  • 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, default = 0.0) – 当 > 0.0 时，在残差块的主路径中对每个样本应用随机深度。

  • fuse_qkv_params (bool, default = ‘False’) – 如果设置为 True，则 TransformerLayer 模块为 query-key-value 公开单个融合参数。 这可以实现诸如无连接/拆分的 QKV 融合之类的优化，并且还启用参数 fuse_wgrad_accumulation。

forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, self_attn_mask_type: str | None = None, window_size: Tuple[int, int] | None = None, encoder_output: torch.Tensor | None = None, enc_dec_attn_mask: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None, enc_dec_attn_mask_type: str | None = None, enc_dec_window_size: Tuple[int, int] | None = None, is_first_microbatch: bool | None = None, checkpoint_core_attention: bool = False, inference_params: transformer_engine.pytorch.attention.InferenceParams | None = None, rotary_pos_emb: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None, core_attention_bias_type: str = 'no_bias', core_attention_bias: torch.Tensor | None = None, alibi_slopes: torch.Tensor | None = None, cu_seqlens_q: torch.Tensor | None = None, cu_seqlens_kv: torch.Tensor | None = None, max_seqlen_q: int | None = None, max_seqlen_kv: int | None = None, fast_zero_fill: bool = True) → torch.Tensor

Transformer Layer：attention 块和前馈网络 (MLP)

注意

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

参数:

  • hidden_states (torch.Tensor) – 输入张量。

  • attention_mask (Optional[torch.Tensor], default = None) – 布尔张量，用于 mask out 自注意力 softmax 输入。 对于 padding mask，它应为 [batch_size, 1, 1, seqlen_q]，对于 “arbitrary” mask，它应可广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]。 对于因果 mask 和 “no_mask” 类型，它应为 None。 True 值表示相应位置被 mask out，False 值表示允许该位置参与 attention。

  • self_attn_mask_type ({'no_mask', 'causal', 'padding', 'padding_causal',) – ‘causal_bottom_right’, ‘padding_causal_bottom_right’,’arbitrary’}, default = causal 传递到编码器的 softmax 操作中的 attention mask 类型。 默认情况下，因果 mask 与 softmax 矩阵的左上角对齐。 当 mask 类型中指定 “bottom_right” 时，因果 mask 与右下角对齐。

  • window_size (Optional[Tuple[int, int]], default = None) – 编码器中本地注意力的滑动窗口大小。

  • encoder_output (Optional[torch.Tensor], default = None) – 编码器块的输出，如果使用 layer_type=”decoder”，则馈送到解码器块。

  • enc_dec_attn_mask (Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]],) – default = None。 布尔张量，如果使用 layer_type=”decoder”，则用于 mask out 互注意力 softmax 输入。 对于 padding mask，它应为 [batch_size, 1, 1, seqlen_q] 和 [batch_size, 1, 1, seqlen_kv] 的两个 mask 的元组。 对于 “arbitrary” mask，它应可广播到 [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]。 对于因果 mask 和 “no_mask”，它应为 None。 True 值表示相应位置被 mask out，False 值表示允许该位置参与 attention。

  • enc_dec_attn_mask_type ({'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'},) – default = None 传递到解码器的 softmax 操作中的 attention mask 类型。

  • enc_dec_window_size (Optional[Tuple[int, int]], default = None) – 解码器中本地注意力的滑动窗口大小。

  • is_first_microbatch ({True, False, None}, default = None) –

    在使用梯度累积或流水线并行进行训练期间，数据的小批量会进一步拆分为微批量。在同一小批量的微批量之间，模型权重不会更新。设置此参数指示当前微批量是否为小批量中的第一个。设置后，此参数启用其他优化

    • 在 FP8 训练期间，它允许缓存权重的 FP8 版本

    • 它还允许跳过第一个微批量期间的梯度累积（因为它是生成的第一个梯度）

  • checkpoint_core_attention (bool, default = False) – 如果为 true，则在反向传播期间重新计算核心注意力的前向激活，以节省内存，否则这些内存将被占用以存储前向激活直到反向传播。

  • rotary_pos_emb (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = None) – 用于 query 和 key 张量的嵌入，以应用旋转位置嵌入。 默认情况下，不应用输入嵌入。

  • core_attention_bias_type (str, default = no_bias) – 偏置类型，{no_bias, pre_scale_bias, post_scale_bias, alibi}

  • core_attention_bias (Optional[torch.Tensor], default = None) – Q * K.T 的偏置张量

  • alibi_slopes (Optional[torch.Tensor], default = None) – FP32 中的 ALiBi 斜率，形状为 [nheads] 或 [batch_size, nheads]。它将 (-alibi_slope * (i + seqlen_k - seqlen_q - j)) 的偏置添加到查询 i 和键 j 的注意力分数中。

  • cu_seqlens_q (Optional[torch.Tensor], default = None) – 批次中 query_layer 的序列长度的累积和（无偏移量），形状为 [batch_size + 1]，dtype 为 torch.int32。

  • cu_seqlens_kv (Optional[torch.Tensor], default = None) – 批次中 key_layer 和 value_layer 的序列长度的累积和（无偏移量），形状为 [batch_size + 1]，dtype 为 torch.int32。

  • max_seqlen_q (Optional[int], default = None) – query_layer 中的最大序列长度。 如果未提供，则从 cu_seqlens_q 计算。

  • max_seqlen_kv (Optional[int], default = None) – key_layer 和 value_layer 中的最大序列长度。 如果未提供，则从 cu_seqlens_kv 计算。

  • fast_zero_fill (bool, default = True) – 是否在使用前将输出张量设置为 0。

  • inference_params (InferenceParams, default = None) – 传递到主模型的推理参数，以便在推理期间有效地计算和存储上下文。

set_context_parallel_group(cp_group: transformer_engine.pytorch.constants.dist_group_type | List[transformer_engine.pytorch.constants.dist_group_type] | None, cp_global_ranks: List[int], cp_stream: torch.cuda.Stream, cp_comm_type: str = 'p2p') → None

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

参数:

  • cp_group (Union[ProcessGroup, List[ProcessGroup]]) – 上下文并行进程组。ProcessGroup 用于 “p2p”、“all_gather” 和 “a2a” 的 cp_comm_type。List[ProcessGroup] 用于 “a2a+p2p” 的 cp_comm_type，其中 cp_group[0] 和 cp_group[1] 分别用于 a2a 和 p2p 通信。

  • cp_global_ranks (List[int]) – 上下文组中的全局 ranks 列表。

  • cp_stream (torch.cuda.Stream) – 用于上下文并行执行的 cuda 流。

  • cp_comm_type (str, default = p2p) –

    上下文并行性的GPU间通信类型。可以是 “p2p”、“all_gather”、“a2a” 或 “a2a+p2p”。“p2p”：在环形拓扑中使用 P2P 通信交换 KV 块。

    P2P 是异步的，可以与注意力计算重叠。

    ”all_gather”：在 attention 之前，使用 all-gather 获取完整的 KV 序列。

    all-gather 不是异步的，不能重叠。

    ”a2a”：类似于 DeepSpeed Ulysses，在 CP 之间分散 attention heads

    组中，并进行 gather 以获取完整的 QKV 序列。

    ”a2a+p2p”：分层 CP 实现。首先在每个 CP 子组内对 QKV 应用 a2a（例如，通过 NVLink），然后在子组之间使用 p2p 交换 KV（例如，通过 IBLink）。

set_tensor_parallel_group(tp_group: transformer_engine.pytorch.constants.dist_group_type | None) → 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) → None

添加新的 RNG 状态。

name: str

RNG 状态的字符串标识符。

seed: int

RNG 状态的 PyTorch 种子。

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

Fork CUDA rng 状态，执行操作，然后以原始状态退出。

name: str

RNG 状态的字符串标识符。

get_states() → Dict[str, torch.Tensor]

获取 rng 状态。复制字典，以便我们拥有指向状态的直接指针，而不仅仅是指向字典的指针。

reset()

设置为初始状态（无跟踪器）。

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

设置 rng 状态。出于效率考虑，我们不检查种子的兼容性大小。

states: Dict[str, torch.Tensor]

从字符串名称到 RNG 状态的映射。

transformer_engine.pytorch.fp8_autocast(enabled: bool = True, calibrating: bool = False, fp8_recipe: transformer_engine.common.recipe.Recipe | None = None, fp8_group: transformer_engine.pytorch.constants.dist_group_type | None = None, _graph: bool = False) → 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 和 scale 数据。这对于在使用更高精度进行训练时保存可用于推理的 fp8 检查点非常有用。

  • fp8_recipe (recipe.Recipe, default = None) – 用于 FP8 训练的 recipe。

  • fp8_group (torch._C._distributed_c10d.ProcessGroup, default = None) – 在每个训练步骤结束时缩减 fp8 张量的 amax 的分布式组。

transformer_engine.pytorch.fp8_model_init(enabled: bool = True, recipe: transformer_engine.common.recipe.Recipe | None = None) → None

用于参数 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 类似微调，其中模型的主要参数不会更改。

  • recipe (transformer_engine.common.recipe.Recipe, default = None) –

    用于创建参数的 Recipe。如果保留为 None，则使用默认的 FP8 recipe。

    此功能是实验性的。

transformer_engine.pytorch.checkpoint(function: Callable, *args: Tuple[torch.Tensor, Ellipsis], **kwargs: Dict[str, Any]) → Tuple[torch.Tensor, Ellipsis]

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

警告

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

警告

use_reentrant=False 不支持提前停止，并且在后向传递中重新计算激活时，将执行检查点模块的整个前向传递。

参数:

  • function (Callable) – 用于使用指定的 args 和 kwargs 运行前向和后向传递的 pytorch 模块。

  • distribute_saved_activations (bool, default = False) – 如果设置为 True 且 use_reentrant=True，则第一个张量参数在指定的张量并行组 (tp_group) 中分发，然后再保存以用于后向传递。当 use_reentrant=False 时，此操作无效。

  • get_rng_state_tracker (Callable, default = None) – 返回 CudaRNGStatesTracker() 实例的 python 可调用对象。

  • tp_group (ProcessGroup, default = None) – 张量并行进程组。仅当 distribute_saved_activations=True 且 use_reentrant=True 时使用。如果为 None，则回退到默认组。

  • use_reentrant (bool, default = True) – 在可重入模式下执行检查点。

  • args (tuple) – 用于 function 输入的 torch 张量元组。

  • kwargs (dict) – 用于 function 关键字参数的字符串键字典。

transformer_engine.pytorch.make_graphed_callables(modules: SingleOrTuple[Callable], sample_args: SingleOrTuple[Tuple[torch.Tensor, Ellipsis]], num_warmup_iters: int = 3, allow_unused_input: bool = False, sample_kwargs: SingleOrTuple[Dict[str, Any]] | None = None, fp8_enabled: bool = False, fp8_calibrating: bool = False, fp8_recipe: transformer_engine.common.recipe.DelayedScaling | None = None, fp8_group: transformer_engine.pytorch.constants.dist_group_type | None = None, fp8_weight_caching: bool = False, _order: List[int] | None = None, pool: Tuple[int, Ellipsis] | None = None, retain_graph_in_backward: bool = False) → Callable | Tuple[Callable, Ellipsis]

制作 Transformer Engine 模块的 CUDA 图版本

PyTorch 的 make_graphed_callables 实用函数的变体，支持 Transformer Engine 模块和 FP8。有关更多文档，请参阅 原始 PyTorch 实现。

图参数:

  • modules (((元组) 的) 可调用对象) – 要制作图的可调用对象或可调用对象元组。

  • sample_args (((元组) 的) torch.Tensor 元组) – 可调用对象的位置参数。

  • num_warmup_iters (int, default = 3) – 预热迭代次数。

  • allow_unused_input (bool, default = False) – 是否处理可调用对象的输入和输出在计算图中断开连接的情况。

  • sample_kwargs (((元组) 的) dict, optional) – 可调用对象的关键字参数

  • pool ((元组) int, default = None, optional) – 从函数 torch.cuda.graph_pool_handle 返回的实例，提示此图可能与指示的池共享内存。

  • retain_graph_in_backward (bool, default = False) – 是否在后向图捕获中设置 retain_graph=True。

FP8 相关参数:

  • fp8_enabled (bool, default = True) – 是否启用 fp8

  • fp8_calibrating (bool, default = False) – 即使在未启用 fp8 的情况下执行，校准模式也允许收集 fp8 张量的统计信息，例如 amax 和 scale 数据。这对于在使用更高精度进行训练时保存可用于推理的 fp8 检查点非常有用。

  • fp8_recipe (recipe.DelayedScaling, default = None) – 用于 FP8 训练的 recipe。

  • fp8_group (torch._C._distributed_c10d.ProcessGroup, default = None) – 在每个训练步骤结束时缩减 fp8 张量的 amax 的分布式组。

  • fp8_weight_caching (bool, default = False) – 是否跨微批次缓存 FP8 权重。如果设置为 True，则必须将 is_first_microbatch 布尔参数传递到 TransformerEngine 模块的前向方法中。当使用 TE 的 fp8_model_init API 在 FP8 中存储主权重并使用 FP8 感知优化器时，如果在 TE 外部（例如，在优化器步骤中）计算权重转置，则此参数必须设置为 False。

transformer_engine.pytorch.get_cpu_offload_context(enabled: bool = False, num_layers: int = 1, model_layers: int = 1, offload_activations: bool = True, offload_weights: bool = True)

此函数返回 CPU 卸载上下文和同步器函数，该函数需要在每个 Transformer 层之后使用。如果未启用卸载，则返回 nullcontext()。

用法

cpu_offload_context, cpu_offload_synchronizer = get_cpu_offload_context(enabled=True)

with cpu_offload_context:
    te_layer.forward(inp_tensor)
cpu_offload_synchronizer()
参数:

  • enabled (bool, default = False) – 设置为 True 时，启用 CPU 卸载功能。

  • num_layers (int, default = 1) – 确定要卸载激活/权重的 Transformer 层数。

  • model_layers (int, default = 1) – 在此上下文下使用的模型层数。

  • offload_activations (bool, default = True) – 设置为 True 时，卸载 TE 层的激活。

  • offload_weights (bool, default = True) – 设置为 True 时，卸载 TE 层的权重。

transformer_engine.pytorch.moe_permute(inp: torch.Tensor, indices: torch.Tensor, num_out_tokens: int = -1, max_token_num: int = -1) → Tuple[torch.Tensor, torch.Tensor]

根据索引置换令牌。具有相同索引的令牌将被分组在一起。

参数:

  • inp (torch.Tensor) – 形状为 [num_tokens, hidden_size] 的输入张量，将在其上应用置换。

  • indices (torch.Tensor) – 形状为 [num_tokens, topK] 且 dtype 为 ‘int32’ 的令牌到专家索引张量。

  • num_out_tokens (int, default = -1) – 有效输出令牌计数，表示未删除的令牌数。默认情况下，设置为 '-1'，表示不删除任何令牌。

  • max_token_num (int, default = -1) – 令牌的最大数量，用于工作区分配。默认情况下，设置为 '-1'，表示工作区大小的计算由运算符自动接管。

transformer_engine.pytorch.moe_unpermute(inp: torch.Tensor, row_id_map: torch.Tensor, probs: torch.Tensor = None) → torch.Tensor

取消置换具有置换令牌的张量，并可选择将令牌与其对应的概率合并。

参数:

  • inp (torch.Tensor) – 具有要取消置换的形状为 [num_tokens, hidden_size] 的置换令牌的输入张量。

  • row_id_map (torch.Tensor) – 用于取消置换令牌的排序索引的映射表张量，它是 Permute 的第二个输出张量。

  • probs (torch.Tensor) – 与置换令牌对应的概率张量。如果提供，则取消置换的令牌将与其各自的概率合并。默认情况下，设置为空张量，这意味着令牌通过累积直接合并。

transformer_engine.pytorch.initialize_ub(shape: list, tp_size: int, use_fp8: bool = False, dtype: torch.dtype = torch.bfloat16, ub_cfgs: dict | None = None, bootstrap_backend: str | torch.distributed.Backend = None) → None

初始化 Userbuffers 通信器，用于在 te.Linear、te.LayerNormLinear 和 te.LayerNormMLP 模块中将张量并行通信与 GEMM 计算重叠。

参数:

  • shape (list) – 通信缓冲区的形状，通常设置为与 te.TransformerLayer 前向传递的输入张量的全局形状相同，序列和批次维度折叠在一起 - 即：(sequence_length * batch_size, hidden_size)

  • tp_size (int) – 张量并行进程组中的 GPU 数量

  • use_fp8 (bool = False) – 为 FP8 GEMM 输入/输出分配通信缓冲区

  • dtype (torch.dtype = torch.bfloat16) – 当 use_fp8 = False 时，通信缓冲区的非 FP8 数据类型

  • ub_cfgs (dict = None) –

    配置字典，结构为 ``` {

    <gemm_name>{

    “method”: <”ring_exchange” 或 “pipeline”>, “is_reduce_scatter”: bool, “num_sm”: int, “cga_size”: int, “set_sm_margin”: bool, “num_splits”: int, “aggregate”: bool, “atomic_gemm”: bool, “use_ce”: bool, “fp8_buf”: bool,

    }

    对于 te.TransformerLayer GEMM 层，位于 [“qkv_fprop”, “qkv_dgrad”, “qkv_wgrad”, “proj_fprop”, “proj_dgrad”, “proj_wgrad”, “fc1_fprop”, “fc1_dgrad”, “fc2_dgrad”, “fc2_fprop”, “fc2_dgrad”] 中。

  • bootstrap_backend (str = None) – torch.distributed 通信后端，用于 Userbuffers 初始化期间的 all-gather、broadcast 和 barrier 集合通信。并非所有后端都对每个集群配置和分布式启动方法都有效，即使它们在 PyTorch 中可用。如果未设置，则初始化首选使用 MPI 后端，如果 MPI 不可用，则首先回退到 Gloo，然后回退到 NCCL。在构建 TE 时设置 NVTE_UB_WITH_MPI=1 会覆盖此选项，并且始终使用 C++ 中的直接 MPI 调用初始化 Userbuffers，这也需要在编译时设置 MPI_HOME=/path/to/mpi/root。

transformer_engine.pytorch.destroy_ub()

销毁所有已分配的 userbuffer 通信器。