fusions 软件包
此软件包提供常用融合操作的模块。融合操作通过增加每次从内存读取张量时完成的工作量来提高计算效率。为了执行融合，此软件包中的模块要么依赖 PyTorch 功能进行即时编译（即旧版 PyTorch 中的 torch.jit.script 或最新版本中的 torch.compile），要么调用外部库（如 Apex 或 TransformerEngine）中的自定义内核。
此模块使用 PyTorch JIT 融合偏置加法和 dropout 操作。由于 dropout 在推理期间不使用，因此在训练模式和推理模式下使用不同的函数。
- core.fusions.fused_bias_dropout.bias_dropout_add_fused_inference(x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float) → torch.Tensor
- core.fusions.fused_bias_dropout.bias_dropout_add_fused_train(x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float) → torch.Tensor
- core.fusions.fused_bias_dropout.bias_dropout_add_unfused(training)
- core.fusions.fused_bias_dropout.get_bias_dropout_add(training, fused)
此模块使用 PyTorch JIT 融合偏置加法和 GeLU 非线性运算。
- class core.fusions.fused_bias_gelu.GeLUFunction(*args: Any, **kwargs: Any)
基类：
torch.autograd.Function- classmethod apply(*args, **kwargs)
- static backward(ctx, grad_output)
- static forward(ctx, input, bias)
- core.fusions.fused_bias_gelu.bias_gelu(bias, y)
- core.fusions.fused_bias_gelu.bias_gelu_back(g, bias, y)
- core.fusions.fused_bias_gelu.bias_gelu_impl(*args, **kwargs)
此模块提供了 Apex 中各种融合 LayerNorm 实现的包装器。
- class core.fusions.fused_layer_norm.FusedLayerNorm(*args: Any, **kwargs: Any)
基类：
torch.nn.ModuleLayer Norm，融合到单个 CUDA 内核中。
- 参数
hidden_size (int) – Transformer 隐藏维度。
eps (float) – 添加到分母的 epsilon，用于数值稳定性。
persist_layer_norm (bool) – 使用持久化融合层归一化内核。
请注意 (此内核仅支持一组隐藏大小。) –
支持。 (如果您的隐藏大小受支持，请检查 persist_ln_hidden_sizes) –
zero_centered_gamma (bool) – 调整 LayerNorm 权重，使其
稳定性。 (以零为中心。这提高了数值) –
config (TransformerConfig) – Transformer 配置。包含以匹配自定义
接口。 (层归一化) –
normalization (str) – 归一化类型，用于 Transformer Engine。
此处。 (必须等于 'LayerNorm') –
- forward(input: torch.Tensor) → torch.Tensor
- reset_parameters()
此模块提供 Apex 中 Softmax 变体的包装器。
- class core.fusions.fused_softmax.FusedScaleMaskSoftmax(*args: Any, **kwargs: Any)
基类：
torch.nn.Module融合操作：缩放 + 掩码 + softmax
- 参数
input_in_fp16 – 标志，指示输入是否为 fp16 数据格式。
input_in_bf16 – 标志，指示输入是否为 bf16 数据格式。
attn_mask_type – 注意力掩码类型（pad 或 causal）
scaled_masked_softmax_fusion – 标志，指示用户想要使用 softmax 融合
mask_func – 要应用的掩码函数。
softmax_in_fp32 – 如果为 true，则以 fp32 精度执行 softmax。
scale – 输入张量缩放中使用的缩放因子。
- forward(input: torch.Tensor, mask: Optional[torch.Tensor])
带掩码输入的 softmax 前向传播。
如果 attn_mask_type 是 causal，则会生成掩码，并且可以传递 None。仅当 attn_mask_type 不是 causal 时，才需要用户定义的掩码。
- forward_fused_softmax(input, mask)
- forward_torch_softmax(input, mask)
- static get_batch_per_block(sq, sk, b, np)
- is_kernel_available(mask, b, np, sq, sk)
- class core.fusions.fused_softmax.ScaledMaskedSoftmax(*args: Any, **kwargs: Any)
基类：
torch.autograd.Function融合操作，按顺序执行以下三个操作：1. 缩放张量。2. 应用掩码。3. 执行 softmax。
- static backward(ctx, output_grads)
- static forward(ctx, inputs, mask, scale)
- class core.fusions.fused_softmax.ScaledSoftmax(*args: Any, **kwargs: Any)
基类：
torch.autograd.Function融合操作，按顺序执行以下两个操作：1. 缩放张量。2. 执行 softmax。
- static backward(ctx, output_grads)
- static forward(ctx, inputs, scale)
- class core.fusions.fused_softmax.ScaledUpperTriangMaskedSoftmax(*args: Any, **kwargs: Any)
基类：
torch.autograd.Function融合操作，按顺序执行以下三个操作：1. 缩放张量。2. 应用上三角掩码（通常用于 gpt 模型）。3. 执行 softmax。
- static backward(ctx, output_grads)
- static forward(ctx, inputs, scale)
此模块使用 PyTorch JIT 融合交叉熵损失计算和批处理通信调用。
- core.fusions.fused_cross_entropy.calculate_cross_entropy_loss(exp_logits: torch.Tensor, predicted_logits_sum_exp_logits: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]
- core.fusions.fused_cross_entropy.calculate_gradients(softmax: torch.Tensor, grad_output: torch.Tensor, target_mask: torch.Tensor, masked_target_1d: torch.Tensor) → torch.Tensor
- core.fusions.fused_cross_entropy.calculate_logits_max(vocab_parallel_logits: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]
- core.fusions.fused_cross_entropy.calculate_predicted_logits(vocab_parallel_logits: torch.Tensor, target: torch.Tensor, logits_max: torch.Tensor, vocab_start_index: int, vocab_end_index: int) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
- core.fusions.fused_cross_entropy.fused_vocab_parallel_cross_entropy(vocab_parallel_logits, target)
当 logits 跨张量并行 ranks 分割时执行交叉熵损失
- 参数
vocab_parallel_logits – 跨张量并行 ranks 分割的 logits，维度为 [sequence_length, batch_size, hidden_size]
target – 正确的词汇表 id，维度为 [sequence_length, micro_batch_size]