class core.tensor_parallel.cross_entropy.VocabParallelCrossEntropy

基类：object

计算交叉熵损失，在张量并行 ranks 之间分割词汇表大小。此实现在融合和非融合交叉熵实现中均有使用

static calculate_cross_entropy_loss(exp_logits: torch.Tensor, predicted_logits: torch.Tensor, sum_exp_logits: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]: 计算交叉熵损失。

static calculate_gradients(grad_2d: torch.Tensor, arange_1d: torch.Tensor, masked_target_1d: torch.Tensor, softmax_update: torch.Tensor, grad_input: torch.Tensor, grad_output: torch.Tensor) → torch.Tensor: 计算梯度。

static calculate_logits_max(vocab_parallel_logits: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]: 计算 logits_max。

static 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, torch.Tensor]: 计算预测 logits。

static prepare_gradient_calculation_operands(softmax: torch.Tensor, target_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 准备梯度计算操作数。

core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy(vocab_parallel_logits, target, label_smoothing=0.0)

当 logits 在张量并行 ranks 之间分割时执行交叉熵损失

参数

vocab_parallel_logits – 在张量并行 ranks 之间分割的 logits，维度为 [sequence_length, batch_size, vocab_size/num_parallel_ranks]
target – 正确的词汇表 ID，维度为 [sequence_length, micro_batch_size]
label_smoothing – 平滑因子，必须在范围 [0.0, 1.0) 内，默认值是无平滑 (=0.0)

tensor_parallel.data 模块

core.tensor_parallel.data.broadcast_data(keys, data, datatype)

将数据从每个模型并行组的 rank 零广播到同一模型并行组的成员。

参数

keys – 要广播的数据字典中的键列表
data – 数据字典，包含字符串键和 cpu 张量值。
datatype – 与键关联的数据中所有张量的 torch 数据类型。

tensor_parallel.layers 模块

class core.tensor_parallel.layers.ColumnParallelLinear(*args: Any, **kwargs: Any)

基类：torch.nn.Module

具有列并行性的线性层。

线性层定义为 Y = XA + b。A 沿其第二个维度并行化为 A = [A_1, …, A_p]。

参数

input_size – 矩阵 A 的第一个维度。
output_size – 矩阵 A 的第二个维度。
bias – 如果为 true，则添加偏置
gather_output – 如果为 true，则对输出调用 all-gather 并使所有 GPU 均可使用 Y，否则，每个 GPU 将具有其输出 Y_i = XA_i
init_method – 初始化权重的方法。请注意，偏置始终设置为零。
stride – 用于跨步线性层。
keep_master_weight_for_test – 这是为测试添加的，应设置为 False。它返回用于初始化的主权重。
skip_bias_add – 如果为 True，则不添加偏置项，而是将其返回以供调用方添加。这可以实现性能优化，其中偏置可以与其他逐元素操作融合。
skip_weight_param_allocation – 如果为 True，则不分配权重参数，并且必须在正向传递期间作为关键字参数 weight 传递。请注意，这不会影响偏置，如果 bias 为 True，则会分配偏置。默认为 False。
embedding_activation_buffer – 当 defer_embedding_wgrad_compute 启用时，此缓冲区保存最后一个流水线阶段的最终嵌入线性层的输入激活。
grad_output_buffer – 当 defer_embedding_wgrad_compute 启用时，此缓冲区保存最后一个流水线阶段的最终嵌入线性层的梯度输出。
is_expert – 如果为 True，则该层被视为 MoE 专家层。
config – ModelParallelConfig 对象
tp_comm_buffer_name – 通信缓冲区名称未在非 Transformer-Engine 模块中使用。
disable_grad_reduce – 如果为 True，则将禁用跨张量并行 ranks 的输出梯度缩减。默认为 False。Nemo 中的 Lora 适配器使用此功能来延迟和融合缩减以及其他梯度，以进行性能优化。

forward(input_: torch.Tensor, weight: Optional[torch.Tensor] = None, runtime_gather_output: Optional[bool] = None)

ColumnParallelLinear 的正向传播

参数

input – 3D 张量，维度顺序为 [sequence, batch, hidden]
weight (可选) – 要使用的权重张量，当 skip_weight_param_allocation 为 True 时是必需的。
runtime_gather_output (bool) – 在运行时收集输出。默认值 None 表示将使用构造函数中的 gather_output 参数。

返回

output
bias

get_extra_state() → None: 保持与 TE 状态字典的兼容性。

set_extra_state(state: Any): 额外状态被忽略

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None): 沿轴 0 分片，偏置分片

class core.tensor_parallel.layers.LinearWithFrozenWeight(*args: Any, **kwargs: Any)

基类：torch.autograd.Function

不计算权重梯度的线性运算符。此操作和 LinearWithGradAccumulationAndAsyncCommunication 在数学上执行相同的正向传播和 DGRAD。

从概念上讲，此操作与 torch.nn.functional.linear 与 weight.requires_grad==False 相同，但在实验中，它们在数学上并不完全相同。

static backward(ctx, grad_output): 使用冻结权重进行反向传播。

static forward(ctx, input, weight, bias, allreduce_dgrad): 使用冻结权重进行正向传播。

class core.tensor_parallel.layers.LinearWithGradAccumulationAndAsyncCommunication(*args: Any, **kwargs: Any)

基类：torch.autograd.Function

请参阅 linear_with_grad_accumulation_and_async_allreduce

static backward(ctx, grad_output): 反向传播。

static forward(ctx, input, weight, bias, gradient_accumulation_fusion, allreduce_dgrad, sequence_parallel, grad_output_buffer, wgrad_deferral_limit): 正向传播。

class core.tensor_parallel.layers.RowParallelLinear(*args: Any, **kwargs: Any)

基类：torch.nn.Module

具有行并行性的线性层。

线性层定义为 Y = XA + b。A 沿其第一个维度并行化，X 沿其第二个维度并行化。A = transpose([A_1 .. A_p]) X = [X_1, …, X_p]

参数

input_size – 矩阵 A 的第一个维度。
output_size – 矩阵 A 的第二个维度。
bias – 如果为 true，则添加偏置。请注意，偏置未并行化。
input_is_parallel – 如果为 true，我们假设输入已在 GPU 之间分割，并且我们不再分割。
init_method – 初始化权重的方法。请注意，偏置始终设置为零。
stride – 用于跨步线性层。
keep_master_weight_for_test – 这是为测试添加的，应设置为 False。它返回用于初始化的主权重。
skip_bias_add – 如果为 True，则不添加偏置项，而是将其返回以供调用方添加。这可以实现性能优化，其中偏置可以与其他逐元素操作融合。
is_expert – 如果为 True，则该层被视为 MoE 专家层
tp_comm_buffer_name – 通信缓冲区名称。未在非 Transformer-Engine 模块中使用。
config – ModelParallelConfig 对象

forward(input_)

RowParallelLinear 的正向传播

参数

input – 3D 张量，维度顺序为 [sequence, batch, hidden]

返回

output
bias

get_extra_state() → None: 保持与 TE 状态字典的兼容性。

set_extra_state(state: Any): 额外状态被忽略

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None): 沿轴 1 分片，偏置未分片

class core.tensor_parallel.layers.VocabParallelEmbedding(*args: Any, **kwargs: Any)

基类：torch.nn.Module

在词汇表维度中并行化的嵌入。

这主要改编自 torch.nn.Embedding，并保留了所有默认值。

参数

num_embeddings – 词汇表大小。
embedding_dim – 隐藏状态的大小。
reduce_scatter_embeddings – 决定是否在嵌入查找后执行 ReduceScatter

关键字参数

config – megatron.core.ModelParallelConfig 对象

forward(input_)

正向传播。

参数: input (torch.Tensor) – 输入张量。

sharded_state_dict(prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int]] = (), metadata: Optional[dict] = None) → Dict[str, Any]: 由于 allow_shape_mismatch 参数，嵌入的非默认实现

core.tensor_parallel.layers.copy_tensor_model_parallel_attributes(destination_tensor, source_tensor): 将模型并行属性从一个张量复制到另一个张量。

core.tensor_parallel.layers.linear_with_frozen_weight(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], gradient_accumulation_fusion: bool, allreduce_dgrad: bool, sequence_parallel: bool, grad_output_buffer: Optional[List[torch.Tensor]] = None, wgrad_deferral_limit: None = None, async_grad_allreduce: Optional[bool] = None) → torch.Tensor

执行 weight.requires_grad == False 的线性层。

此函数处理权重冻结（不可训练）的线性层。在正向传播中，它仅保存权重，而不保存输入激活。在反向传播中，它不执行权重梯度计算或权重梯度 allreduce。

参数

input (torch.Tensor 必需): 类似于 torch.nn.functional.linear 的输入

weight (torch.Tensor 必需): 类似于 torch.nn.functional.linear 的权重

bias (torch.Tensor 可选): 类似于 torch.nn.functional.linear 的偏置

gradient_accumulation_fusion (bool 必需): 虚拟参数，用于保持所有正向传播实现函数之间的 API 统一。

allreduce_dgrad (bool, 必需): 执行输入梯度的 allreduce。: 在此，异步和同步 allreduce 是相同的。如果 sequence_parallel 为 True，则此项必须为 False，因为不执行 all reduce。
sequence_parallel (bool 必需): 指示序列: 使用并行性，因此在正向传播中，输入是全部收集的，而在反向传播中，输入梯度是 reduce scattered 的。

grad_output_buffer (List[torch.Tensor] 可选): 虚拟参数，用于保持所有正向传播实现函数之间的 API 统一。

wgrad_deferral_limit (int 可选): 虚拟参数，用于保持所有正向传播实现函数之间的 API 统一。

async_grad_allreduce (bool 可选): 将在 0.11.0 版本中删除。: 请改用 allreduce_dgrad。

core.tensor_parallel.layers.linear_with_grad_accumulation_and_async_allreduce(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], gradient_accumulation_fusion: bool, allreduce_dgrad: bool, sequence_parallel: bool, grad_output_buffer: Optional[List[torch.Tensor]] = None, wgrad_deferral_limit: Optional[int] = 0, async_grad_allreduce: Optional[bool] = None) → torch.Tensor

在线性层执行中使用异步通信和反向传播中的梯度累积融合。

此选项可以选择将反向传播计算的结果累积到现有的梯度缓冲区中，从而避免在梯度计算后执行额外的加法内核。

此外，输入梯度的张量并行 all reduce 可以与权重梯度的计算异步完成。

在序列并行的情况下，输入梯度的 reduce scatter 与权重梯度的计算异步完成。

使用此模块需要环境变量 CUDA_DEVICE_MAX_CONNECTIONS=1。代码中注明了一些集体操作，应在计算内核之前调度这些操作，以使通信与计算重叠，这对于加速是必要的，但对于正确性不是必要的，因此调度程序不会强制排序。设置 CUDA_DEVICE_MAX_CONNECTIONS=1 会强制内核按调用顺序调度。

参数

input (torch.Tensor 必需) – 类似于 torch.nn.functional.linear 的输入
weight (torch.Tensor 必需) – 类似于 torch.nn.functional.linear 的权重
bias (torch.Tensor 可选) – 类似于 torch.nn.functional.linear 的偏置
gradient_accumulation_fusion (bool 必需) – 执行梯度累积融合，需要自定义 CUDA 扩展 fused_weight_gradient_mlp_cuda 模块。要使用 gradient_accumulation_fusion，您必须使用 –cpp_ext 和 –cuda_ext 安装 APEX。例如：“pip install –global-option=”–cpp_ext” –global-option=”–cuda_ext .” ” 请注意，该扩展需要 CUDA>=11。否则，您必须关闭梯度累积融合。”
allreduce_dgrad (bool 必需) – 执行输入梯度的 allreduce。allreduce 与权重梯度的计算异步完成。如果 sequence_parallel 为 True，则此项必须为 False，因为不执行 all reduce。
sequence_parallel (bool 必需) – 指示使用序列并行性，因此在正向传播中输入是全部收集的，而在反向传播中输入梯度是 reduce scattered 的。
grad_output_buffer (List[torch.Tensor] 可选) – 当嵌入表 wgrad 计算被延迟时，用于保存输出梯度的缓冲区。默认为 None。
wgrad_deferral_limit (int 可选) – 应该延迟嵌入权重梯度 GEMM 的微批次数量的限制。通过将其设置为 0 来禁用。默认为 0。
async_grad_allreduce (bool 可选) – 将在 0.11.0 版本中删除。请改用 allreduce_dgrad。

core.tensor_parallel.layers.param_is_not_tensor_parallel_duplicate(param): 如果传入的参数不是另一个 TP rank 上的重复参数，则返回 true。

core.tensor_parallel.layers.set_defaults_if_not_set_tensor_model_parallel_attributes(tensor): 如果尚未显式设置，则设置默认模型并行属性。

core.tensor_parallel.layers.set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride): 将 tp 属性设置为张量

tensor_parallel.mappings 模块

core.tensor_parallel.mappings.all_gather_last_dim_from_tensor_parallel_region(input_): autograd 函数的包装器：正向传播：AG，反向传播 RS <last dim>

core.tensor_parallel.mappings.all_to_all(group, input_, output_split_sizes_=None, input_split_sizes=None): autograd 函数的包装器

core.tensor_parallel.mappings.all_to_all_hp2sp(input_)

在张量并行组上执行 AlltoAll 通信，将输入张量从形状 [num_tokens, H/TP] 转换为 [num_tokens/TP, H]。

参数: input (torch.Tensor) – 已沿隐藏维度分布的输入张量。
返回: 形状为 [num_tokens/TP, H] 的输出张量。
返回类型: torch.Tensor

core.tensor_parallel.mappings.all_to_all_sp2hp(input_)

在张量并行组上执行 AlltoAll 通信，将输入张量从形状 [num_tokens/TP, H] 转换为 [num_tokens, H/TP]。

参数: input (torch.Tensor) – 已沿序列维度分布的输入张量。
返回: 形状为 [num_tokens, H/TP] 的输出张量。
返回类型: torch.Tensor

core.tensor_parallel.mappings.copy_to_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：copy，反向传播 allreduce

core.tensor_parallel.mappings.gather_from_sequence_parallel_region(input_, tensor_parallel_output_grad=True, group=None, output_split_sizes=None, use_global_buffer=False): autograd 函数的包装器：正向传播：AG，反向传播：RS <first dim>

core.tensor_parallel.mappings.gather_from_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：AG，反向传播：split <last dim>

core.tensor_parallel.mappings.reduce_from_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：all reduce，反向传播 copy

core.tensor_parallel.mappings.reduce_scatter_last_dim_to_tensor_parallel_region(input_): autograd 函数的包装器：正向传播：RS，反向传播 AG：AG <last dim>

core.tensor_parallel.mappings.reduce_scatter_to_sequence_parallel_region(input_, group=None, input_split_sizes=None, use_global_buffer=False): autograd 函数的包装器：正向传播：RS，反向传播 AG <fisrt dim>

core.tensor_parallel.mappings.scatter_to_sequence_parallel_region(input_): autograd 函数的包装器：正向传播：split，反向传播：AG <last dim>

core.tensor_parallel.mappings.scatter_to_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：RS，反向传播：AG <last dim>

tensor_parallel.random 模块

class core.tensor_parallel.random.CheckpointFunction(*args: Any, **kwargs: Any)

基类：torch.autograd.Function

检查点函数

此函数改编自 torch.utils.checkpoint，主要有两处更改：1) torch.cuda.set_rng_state 替换为 _set_cuda_rng_state 2) 模型并行跟踪器中的状态也得到正确跟踪/设置/重置。

static backward(ctx, *args): 反向传播。

static forward(ctx, run_function, distribute_saved_activations, *args): 正向传播。

class core.tensor_parallel.random.CudaRNGStatesTracker

基类：object

cuda RNG 状态的跟踪器。

使用 add 方法，cuda rng 状态基于输入 seed 初始化，并分配给 name。稍后，通过 fork rng 状态，我们可以执行操作并返回到我们的起始 cuda 状态。

add(name, seed): 跟踪 rng 状态。

fork(name='model-parallel-rng'): Fork cuda rng 状态，执行操作，并以原始状态退出。

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

is_initialized(): 检查是否已使用 set_states() 设置内部 RNG 状态。

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

set_states(states): 设置 rng 状态。出于效率目的，我们不检查 seed 的大小以确保兼容性。

core.tensor_parallel.random.checkpoint(function, distribute_saved_activations, *args): 检查点模型或模型的一部分。这直接从 torch.utils.checkpoint 复制而来。

core.tensor_parallel.random.get_cuda_rng_tracker(use_te_rng_tracker=False, inference_rng_tracker=False): 获取 cuda rng 跟踪器。

core.tensor_parallel.random.get_data_parallel_rng_tracker_name(): 获取数据并行 rng 跟踪器名称

core.tensor_parallel.random.get_expert_parallel_rng_tracker_name(): 获取专家并行 rng 跟踪器名称

core.tensor_parallel.random.initialize_rng_tracker(use_te_rng_tracker: bool = False, inference_rng_tracker: bool = False): 创建 RNG 跟踪器。“use_te_rng_tracker”确定是否使用 Megatron 或 TransformerEngine 的实现。特别是，TransformerEngine 的实现是可 CUDA 图化的，并且支持 FP8。

core.tensor_parallel.random.model_parallel_cuda_manual_seed(seed, te_rng_tracker=False, inference_rng_tracker=False)

初始化模型并行 cuda 种子。

应在模型并行初始化后调用此函数。此外，在此函数之后不应调用 torch.cuda.manual_seed。基本上，这是该函数的替代品。跟踪三组 RNG 状态：默认状态：这用于数据并行性，并且在一组模型并行 GPU 之间相同，但在不同的模型并行组之间不同。例如，这用于非张量模型并行区域中的 dropout。张量模型并行状态：此状态在一组模型并行 GPU 之间不同，但在数据并行组之间相同。例如，这用于模型并行区域中的 dropout。专家并行种子：此状态仅用于 MoE 模型的专家层。它在专家张量和专家模型并行 GPU 之间不同，但在专家数据并行组之间相同。

tensor_parallel.utils 模块

class core.tensor_parallel.utils.VocabUtility

基类：object

将词汇表拆分为 world_size 块，并返回属于 rank 分区的词汇表的第一个和最后一个索引：请注意，索引在 [fist, last) 中

static vocab_range_from_global_vocab_size(global_vocab_size: int, rank: int, world_size: int) → Sequence[int]: 从全局词汇表大小计算词汇表范围。

static vocab_range_from_per_partition_vocab_size(per_partition_vocab_size: int, rank, world_size: int) → Sequence[int]: 从每个分区的词汇表大小计算词汇表范围。

core.tensor_parallel.utils.gather_split_1d_tensor(tensor)

split_tensor_into_1d_equal_chunks 的反向操作。从张量模型并行 ranks 收集值。

返回一个新的 Tensor，其中包含收集的数据。

参数: tensor – 张量或此 rank 数据部分的视图。

core.tensor_parallel.utils.split_tensor_along_last_dim(tensor: torch.Tensor, num_partitions: int, contiguous_split_chunks: bool = False) → List[torch.Tensor]

沿张量的最后一个维度分割张量。

参数

tensor – 输入张量。
num_partitions – 分割张量的分区数
contiguous_split_chunks – 如果为 True，则使每个块在内存中连续。

返回

张量列表

core.tensor_parallel.utils.split_tensor_into_1d_equal_chunks(tensor, new_buffer=False)

将张量分割成跨张量并行 rank 的相等 1D 块。

返回一个张量或视图，其中包含此 rank 的数据部分。

参数: tensor – 要分割的张量
关键字参数: new_buffer (bool) – 如果为 True，则返回一个新张量。如果为 False，则返回现有张量的视图。默认为 False

模块内容

class core.tensor_parallel.ColumnParallelLinear(*args: Any, **kwargs: Any)

基类：torch.nn.Module

具有列并行性的线性层。

线性层定义为 Y = XA + b。A 沿其第二个维度并行化为 A = [A_1, …, A_p]。

参数

input_size – 矩阵 A 的第一个维度。
output_size – 矩阵 A 的第二个维度。
bias – 如果为 true，则添加偏置
gather_output – 如果为 true，则对输出调用 all-gather 并使所有 GPU 均可使用 Y，否则，每个 GPU 将具有其输出 Y_i = XA_i
init_method – 初始化权重的方法。请注意，偏置始终设置为零。
stride – 用于跨步线性层。
keep_master_weight_for_test – 这是为测试添加的，应设置为 False。它返回用于初始化的主权重。
skip_bias_add – 如果为 True，则不添加偏置项，而是将其返回以供调用方添加。这可以实现性能优化，其中偏置可以与其他逐元素操作融合。
skip_weight_param_allocation – 如果为 True，则不分配权重参数，并且必须在正向传递期间作为关键字参数 weight 传递。请注意，这不会影响偏置，如果 bias 为 True，则会分配偏置。默认为 False。
embedding_activation_buffer – 当 defer_embedding_wgrad_compute 启用时，此缓冲区保存最后一个流水线阶段的最终嵌入线性层的输入激活。
grad_output_buffer – 当 defer_embedding_wgrad_compute 启用时，此缓冲区保存最后一个流水线阶段的最终嵌入线性层的梯度输出。
is_expert – 如果为 True，则该层被视为 MoE 专家层。
config – ModelParallelConfig 对象
tp_comm_buffer_name – 通信缓冲区名称未在非 Transformer-Engine 模块中使用。
disable_grad_reduce – 如果为 True，则将禁用跨张量并行 ranks 的输出梯度缩减。默认为 False。Nemo 中的 Lora 适配器使用此功能来延迟和融合缩减以及其他梯度，以进行性能优化。

forward(input_: torch.Tensor, weight: Optional[torch.Tensor] = None, runtime_gather_output: Optional[bool] = None)

ColumnParallelLinear 的正向传播

参数

input – 3D 张量，维度顺序为 [sequence, batch, hidden]
weight (可选) – 要使用的权重张量，当 skip_weight_param_allocation 为 True 时是必需的。
runtime_gather_output (bool) – 在运行时收集输出。默认值 None 表示将使用构造函数中的 gather_output 参数。

返回

output
bias

get_extra_state() → None: 保持与 TE 状态字典的兼容性。

set_extra_state(state: Any): 额外状态被忽略

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None): 沿轴 0 分片，偏置分片

class core.tensor_parallel.RowParallelLinear(*args: Any, **kwargs: Any)

基类：torch.nn.Module

具有行并行性的线性层。

线性层定义为 Y = XA + b。A 沿其第一个维度并行化，X 沿其第二个维度并行化。A = transpose([A_1 .. A_p]) X = [X_1, …, X_p]

参数

input_size – 矩阵 A 的第一个维度。
output_size – 矩阵 A 的第二个维度。
bias – 如果为 true，则添加偏置。请注意，偏置未并行化。
input_is_parallel – 如果为 true，我们假设输入已在 GPU 之间分割，并且我们不再分割。
init_method – 初始化权重的方法。请注意，偏置始终设置为零。
stride – 用于跨步线性层。
keep_master_weight_for_test – 这是为测试添加的，应设置为 False。它返回用于初始化的主权重。
skip_bias_add – 如果为 True，则不添加偏置项，而是将其返回以供调用方添加。这可以实现性能优化，其中偏置可以与其他逐元素操作融合。
is_expert – 如果为 True，则该层被视为 MoE 专家层
tp_comm_buffer_name – 通信缓冲区名称。未在非 Transformer-Engine 模块中使用。
config – ModelParallelConfig 对象

forward(input_)

RowParallelLinear 的正向传播

参数

input – 3D 张量，维度顺序为 [sequence, batch, hidden]

返回

output
bias

get_extra_state() → None: 保持与 TE 状态字典的兼容性。

set_extra_state(state: Any): 额外状态被忽略

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None): 沿轴 1 分片，偏置未分片

class core.tensor_parallel.VocabParallelEmbedding(*args: Any, **kwargs: Any)

基类：torch.nn.Module

在词汇表维度中并行化的嵌入。

这主要改编自 torch.nn.Embedding，并保留了所有默认值。

参数

num_embeddings – 词汇表大小。
embedding_dim – 隐藏状态的大小。
reduce_scatter_embeddings – 决定是否在嵌入查找后执行 ReduceScatter

关键字参数

config – megatron.core.ModelParallelConfig 对象

forward(input_)

正向传播。

参数: input (torch.Tensor) – 输入张量。

sharded_state_dict(prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int]] = (), metadata: Optional[dict] = None) → Dict[str, Any]: 由于 allow_shape_mismatch 参数，嵌入的非默认实现

core.tensor_parallel.broadcast_data(keys, data, datatype)

将数据从每个模型并行组的 rank 零广播到同一模型并行组的成员。

参数

keys – 要广播的数据字典中的键列表
data – 数据字典，包含字符串键和 cpu 张量值。
datatype – 与键关联的数据中所有张量的 torch 数据类型。

core.tensor_parallel.checkpoint(function, distribute_saved_activations, *args): 检查点模型或模型的一部分。这直接从 torch.utils.checkpoint 复制而来。

core.tensor_parallel.copy_tensor_model_parallel_attributes(destination_tensor, source_tensor): 将模型并行属性从一个张量复制到另一个张量。

core.tensor_parallel.copy_to_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：copy，反向传播 allreduce

core.tensor_parallel.gather_from_sequence_parallel_region(input_, tensor_parallel_output_grad=True, group=None, output_split_sizes=None, use_global_buffer=False): autograd 函数的包装器：正向传播：AG，反向传播：RS <first dim>

core.tensor_parallel.gather_from_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：AG，反向传播：split <last dim>

core.tensor_parallel.gather_split_1d_tensor(tensor)

split_tensor_into_1d_equal_chunks 的反向操作。从张量模型并行 ranks 收集值。

返回一个新的 Tensor，其中包含收集的数据。

参数: tensor – 张量或此 rank 数据部分的视图。

core.tensor_parallel.get_cuda_rng_tracker(use_te_rng_tracker=False, inference_rng_tracker=False): 获取 cuda rng 跟踪器。

core.tensor_parallel.get_expert_parallel_rng_tracker_name(): 获取专家并行 rng 跟踪器名称

core.tensor_parallel.linear_with_grad_accumulation_and_async_allreduce(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], gradient_accumulation_fusion: bool, allreduce_dgrad: bool, sequence_parallel: bool, grad_output_buffer: Optional[List[torch.Tensor]] = None, wgrad_deferral_limit: Optional[int] = 0, async_grad_allreduce: Optional[bool] = None) → torch.Tensor

在线性层执行中使用异步通信和反向传播中的梯度累积融合。

此选项可以选择将反向传播计算的结果累积到现有的梯度缓冲区中，从而避免在梯度计算后执行额外的加法内核。

此外，输入梯度的张量并行 all reduce 可以与权重梯度的计算异步完成。

在序列并行的情况下，输入梯度的 reduce scatter 与权重梯度的计算异步完成。

参数

input (torch.Tensor 必需) – 类似于 torch.nn.functional.linear 的输入
weight (torch.Tensor 必需) – 类似于 torch.nn.functional.linear 的权重
bias (torch.Tensor 可选) – 类似于 torch.nn.functional.linear 的偏置
gradient_accumulation_fusion (bool 必需) – 执行梯度累积融合，需要自定义 CUDA 扩展 fused_weight_gradient_mlp_cuda 模块。要使用 gradient_accumulation_fusion，您必须使用 –cpp_ext 和 –cuda_ext 安装 APEX。例如：“pip install –global-option=”–cpp_ext” –global-option=”–cuda_ext .” ” 请注意，该扩展需要 CUDA>=11。否则，您必须关闭梯度累积融合。”
allreduce_dgrad (bool 必需) – 执行输入梯度的 allreduce。allreduce 与权重梯度的计算异步完成。如果 sequence_parallel 为 True，则此项必须为 False，因为不执行 all reduce。
sequence_parallel (bool 必需) – 指示使用序列并行性，因此在正向传播中输入是全部收集的，而在反向传播中输入梯度是 reduce scattered 的。
grad_output_buffer (List[torch.Tensor] 可选) – 当嵌入表 wgrad 计算被延迟时，用于保存输出梯度的缓冲区。默认为 None。
wgrad_deferral_limit (int 可选) – 应该延迟嵌入权重梯度 GEMM 的微批次数量的限制。通过将其设置为 0 来禁用。默认为 0。
async_grad_allreduce (bool 可选) – 将在 0.11.0 版本中删除。请改用 allreduce_dgrad。

core.tensor_parallel.model_parallel_cuda_manual_seed(seed, te_rng_tracker=False, inference_rng_tracker=False)

初始化模型并行 cuda 种子。

core.tensor_parallel.param_is_not_tensor_parallel_duplicate(param): 如果传入的参数不是另一个 TP rank 上的重复参数，则返回 true。

core.tensor_parallel.reduce_from_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：all reduce，反向传播 copy

core.tensor_parallel.reduce_scatter_to_sequence_parallel_region(input_, group=None, input_split_sizes=None, use_global_buffer=False): autograd 函数的包装器：正向传播：RS，反向传播 AG <fisrt dim>

core.tensor_parallel.scatter_to_sequence_parallel_region(input_): autograd 函数的包装器：正向传播：split，反向传播：AG <last dim>

core.tensor_parallel.scatter_to_tensor_model_parallel_region(input_): autograd 函数的包装器：正向传播：RS，反向传播：AG <last dim>

core.tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(tensor): 如果尚未显式设置，则设置默认模型并行属性。

core.tensor_parallel.set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride): 将 tp 属性设置为张量

core.tensor_parallel.split_tensor_along_last_dim(tensor: torch.Tensor, num_partitions: int, contiguous_split_chunks: bool = False) → List[torch.Tensor]

沿张量的最后一个维度分割张量。

参数

tensor – 输入张量。
num_partitions – 分割张量的分区数
contiguous_split_chunks – 如果为 True，则使每个块在内存中连续。

返回

张量列表

core.tensor_parallel.split_tensor_into_1d_equal_chunks(tensor, new_buffer=False)

将张量分割成跨张量并行 rank 的相等 1D 块。

返回一个张量或视图，其中包含此 rank 的数据部分。

参数: tensor – 要分割的张量
关键字参数: new_buffer (bool) – 如果为 True，则返回一个新张量。如果为 False，则返回现有张量的视图。默认为 False

core.tensor_parallel.vocab_parallel_cross_entropy(vocab_parallel_logits, target, label_smoothing=0.0)

当 logits 在张量并行 ranks 之间分割时执行交叉熵损失

参数

vocab_parallel_logits – 在张量并行 ranks 之间分割的 logits，维度为 [sequence_length, batch_size, vocab_size/num_parallel_ranks]
target – 正确的词汇表 ID，维度为 [sequence_length, micro_batch_size]
label_smoothing – 平滑因子，必须在范围 [0.0, 1.0) 内，默认值是无平滑 (=0.0)

上一篇 models.bert 包

下一篇 context_parallel 包

tensor_parallel 包

子模块

tensor_parallel.cross_entropy 模块

tensor_parallel.data 模块

tensor_parallel.layers 模块

tensor_parallel.mappings 模块

tensor_parallel.random 模块

tensor_parallel.utils 模块

模块内容