`core.tensor_parallel.layers`#

Module Contents#

Classes#

`VocabParallelEmbedding`	Embedding parallelized in the vocabulary dimension.
`LinearWithFrozenWeight`	Linear operator that does not calculate gradient for weight. This op and LinearWithGradAccumulationAndAsyncCommunication performs mathematically-identical forward and DGRAD.
`LinearWithGradAccumulationAndAsyncCommunication`	See linear_with_grad_accumulation_and_async_allreduce
`ColumnParallelLinear`	Linear layer with column parallelism.
`RowParallelLinear`	Linear layer with row parallelism.

Functions#

`param_is_not_tensor_parallel_duplicate`	Return whether a parameter contributes to a unique model-parallel shard.
`copy_gtp_attributes`	Copy GTP metadata onto a param view/copy (e.g. an optimizer’s master or shard param): dedup tags for `param_is_not_gtp_duplicate`, the checkpoint replica group, and `pad_length`/`group` for `gtp_local_pad_zero_count`. The latter two must both be present or padding exclusion silently returns 0.
`param_is_not_gtp_duplicate`	True if the param’s grad is counted once across the GTP_remat/EGTP_remat axis.
`gtp_local_pad_zero_count`	Count structural GTP alignment-padding elements in `gtp_shard.view(-1)[range_start:range_end]` (see `_gtp_slice_one_param`).
`set_tensor_model_parallel_attributes`	Sets tp attributes to tensor
`set_defaults_if_not_set_tensor_model_parallel_attributes`	Set default model parallel attributes if not set explicitly already.
`copy_tensor_model_parallel_attributes`	Copy model parallel attributes from one tensor to another.
`_initialize_affine_weight_gpu`	Initialize affine weight for model parallel on GPU.
`_initialize_affine_weight_cpu`	Initialize affine weight for model parallel.
`_linear_forward`	Run a linear GEMM with an optional output dtype distinct from its input dtype.
`linear_with_frozen_weight`	Linear layer execution with weight.requires_grad == False.
`_wgrad_gemm`	Weight-gradient GEMM into `out`, which may be wider than the inputs (bf16 -> fp32).
`linear_with_grad_accumulation_and_async_allreduce`	Linear layer execution with asynchronous communication and gradient accumulation fusion in backprop.

Data#

`_grad_accum_fusion_available`
`_MODEL_PARALLEL_ATTRIBUTE_DEFAULTS`

API#

core.tensor_parallel.layers._grad_accum_fusion_available#: True

core.tensor_parallel.layers._MODEL_PARALLEL_ATTRIBUTE_DEFAULTS#: None

core.tensor_parallel.layers.param_is_not_tensor_parallel_duplicate( param, tp_group=None, expert_tp_group=None, )#

Return whether a parameter contributes to a unique model-parallel shard.

Parameters reduced over expert data parallel groups use the expert tensor-parallel group for duplicate filtering. Other parameters use the regular tensor-parallel group.

core.tensor_parallel.layers.copy_gtp_attributes(destination, source)#: Copy GTP metadata onto a param view/copy (e.g. an optimizer’s master or shard param): dedup tags for param_is_not_gtp_duplicate, the checkpoint replica group, and pad_length/group for gtp_local_pad_zero_count. The latter two must both be present or padding exclusion silently returns 0.

core.tensor_parallel.layers.param_is_not_gtp_duplicate(param)#

True if the param’s grad is counted once across the GTP_remat/EGTP_remat axis.

GTP_remat/EGTP_remat shards are unique per peer (kept); replicated params counted only on rank 0 of the gtp_remat/egtp_remat axis (else counted N times). When GTP_remat is off rank is 0, so every param is kept.

core.tensor_parallel.layers.gtp_local_pad_zero_count(gtp_shard, range_start, range_end)#

Count structural GTP alignment-padding elements in gtp_shard.view(-1)[range_start:range_end] (see _gtp_slice_one_param).

Padding is a contiguous suffix of the unsharded padded buffer (shard_dim0 * group.size() rows), sliced evenly across the GTP group. It usually lands entirely in the last rank’s shard, but when pad_length exceeds one shard’s own row count (small dim0 relative to pad_for_alignment * gtp_remat_size) it spills backward from the tail into lower-numbered ranks’ shards too. Computed via each rank’s row offset in the unsharded buffer – not special-cased to the last rank – so both cases come out correct.

Parameters:

gtp_shard – This rank’s local GTP shard (carries pad_length/group).
range_start – Start offset, in gtp_shard.view(-1) flat-index units, of the fragment being queried.
range_end – End offset (exclusive) of that fragment. Pass 0, gtp_shard.numel() for the whole shard (e.g. LayerWiseDistributedOptimizer, which never byte-slices); DistributedOptimizer passes the DP-optimizer-state fragment’s own [param_range.start, param_range.end) instead, since a fragment only covers part of the shard.

core.tensor_parallel.layers.set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride)#: Sets tp attributes to tensor

core.tensor_parallel.layers.set_defaults_if_not_set_tensor_model_parallel_attributes(tensor)#: Set default model parallel attributes if not set explicitly already.

core.tensor_parallel.layers.copy_tensor_model_parallel_attributes( destination_tensor, source_tensor, )#: Copy model parallel attributes from one tensor to another.

core.tensor_parallel.layers._initialize_affine_weight_gpu( weight, init_method, partition_dim, stride=1, is_expert=False, )#: Initialize affine weight for model parallel on GPU.

core.tensor_parallel.layers._initialize_affine_weight_cpu( weight, output_size, input_size, per_partition_size, partition_dim, init_method, stride=1, return_master_weight=False, *, params_dtype=torch.float32, rank=None, world_size=None, skip_set_tensor_parallel_attributes=False, )#

Initialize affine weight for model parallel.

Build the master weight on all processes and scatter the relevant chunk. A weight that is already GTP_remat-sharded is sliced down to this rank’s GTP rows as well, so the initialization matches a run with GTP off.

class core.tensor_parallel.layers.VocabParallelEmbedding( num_embeddings: int, embedding_dim: int, *, init_method: Callable, reduce_scatter_embeddings: bool = False, config: megatron.core.model_parallel_config.ModelParallelConfig, tp_group: Optional[torch.distributed.ProcessGroup] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

Bases: torch.nn.Module

Embedding parallelized in the vocabulary dimension.

This is mainly adapted from torch.nn.Embedding and all the default values are kept.

Parameters:

num_embeddings – vocabulary size.
embedding_dim – size of hidden state.
reduce_scatter_embeddings – Decides whether to perform ReduceScatter after embedding lookup

Keyword Arguments:

config – A megatron.core.ModelParallelConfig object

Initialization

forward(input_)#

Forward.

Parameters:: input_ (torch.Tensor) – Input tensor.

sharded_state_dict( prefix: str = '', sharded_offsets: Tuple[Tuple[int, int, int]] = (), metadata: Optional[dict] = None, ) → core.dist_checkpointing.mapping.ShardedStateDict#: Non-default implementation for embeddings due to allow_shape_mismatch param

class core.tensor_parallel.layers.LinearWithFrozenWeight#

Bases: torch.autograd.Function

Linear operator that does not calculate gradient for weight. This op and LinearWithGradAccumulationAndAsyncCommunication performs mathematically-identical forward and DGRAD.

Conceptually this op is the same as torch.nn.functional.linear with weight.requires_grad==False, but in experiments they are not identical mathematically.

static forward( ctx, input, weight, bias, allreduce_dgrad, tp_group, output_dtype, )#: Forward with frozen weight.

static backward(ctx, grad_output)#: Backward with frozen weight.

core.tensor_parallel.layers._linear_forward( input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], output_dtype: Optional[torch.dtype], ) → torch.Tensor#: Run a linear GEMM with an optional output dtype distinct from its input dtype.

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, tp_group: Optional[torch.distributed.ProcessGroup], grad_output_buffer: Optional[List[torch.Tensor]] = None, wgrad_deferral_limit: None = None, gtp_remat_size: int = 1, output_dtype: Optional[torch.dtype] = None, ) → torch.Tensor#

Linear layer execution with weight.requires_grad == False.

This function handles linear layers with weight frozen (untrainable). In the forward, it only saves weight and does not save input activations. In the backward, it does not perform weight gradient calculation, or weight gradient allreduce.

Args:

input (torch.Tensor required): input like torch.nn.functional.linear

weight (torch.Tensor required): weight like torch.nn.functional.linear

bias (torch.Tensor optional): bias like torch.nn.functional.linear

gradient_accumulation_fusion (bool required): dummy argument, used to keep the API unified between all forward implementation functions.

allreduce_dgrad (bool, required): Do the allreduce of input gradients. Here, async and sync allreduce are the same. If sequence_parallel is True, this must be False, as no all reduce is performed.

sequence_parallel (bool required): Indicates that sequence parallelism is used and thus in the forward pass the input is all gathered, and the backward pass the input gradients are reduce scattered.

tp_group (torch.distributed.ProcessGroup): The process group to use for tensor parallel operations.

grad_output_buffer (List[torch.Tensor] optional): dummy argument, used to keep the API unified between all forward implementation functions.

wgrad_deferral_limit (int optional): dummy argument, used to keep the API unified between all forward implementation functions.

gtp_remat_size (int): GTP shard count. When > 1 the weight is GTP-sharded and must be all-gathered to its full shape before the matmul, mirroring the trainable path. Defaults to 1 (no-op) for the common non-GTP / non-sharded case.

output_dtype (torch.dtype optional): Optional GEMM output dtype. A dtype different from the input dtype requires Transformer Engine general_gemm.

core.tensor_parallel.layers._wgrad_gemm(out, grad_output, total_input)#

Weight-gradient GEMM into out, which may be wider than the inputs (bf16 -> fp32).

Returns out, filled with the weight gradient.

class core.tensor_parallel.layers.LinearWithGradAccumulationAndAsyncCommunication#

Bases: torch.autograd.Function

See linear_with_grad_accumulation_and_async_allreduce

static forward( ctx, input, weight, bias, gradient_accumulation_fusion, allreduce_dgrad, sequence_parallel, grad_output_buffer, wgrad_deferral_limit, tp_group, gtp_remat_size, output_dtype, )#: Forward.

static backward(ctx, grad_output)#: Backward.

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, tp_group: Optional[torch.distributed.ProcessGroup] = None, gtp_remat_size: int = 1, output_dtype: Optional[torch.dtype] = None, ) → torch.Tensor#

Linear layer execution with asynchronous communication and gradient accumulation fusion in backprop.

This has the option to accumulate the result of backprop calculation into an existing gradient buffer, preventing the need to do an additional addition kernel after the gradient calculation.

Additionally, the tensor parallel all reduce of the input gradients can be done asynchronously with the calculation of the weight gradients.

In the case of sequence parallelism, the reduce scatter of the input gradients is done asynchronously with the calculation of the weight gradients.

Use of this module requires that the environment variable CUDA_DEVICE_MAX_CONNECTIONS=1. There are a few collective operations, noted in the code, that should be scheduled before compute kernels to overlap the communication with the computation, which is necessary for a speedup but not for correctness so that ordering isn’t imposed by the scheduler. Setting CUDA_DEVICE_MAX_CONNECTIONS=1 forces the kernels to be scheduled in the order they are called.

Parameters:

input (torch.Tensor required) – input like torch.nn.functional.linear
weight (torch.Tensor required) – weight like torch.nn.functional.linear
bias (torch.Tensor optional) – bias like torch.nn.functional.linear
gradient_accumulation_fusion (bool required) – Perform the gradient accumulation fusion, requires the custom CUDA extension fused_weight_gradient_mlp_cuda module. To use gradient_accumulation_fusion you must install APEX with –cpp_ext and –cuda_ext. For example: “pip install –global-option=”–cpp_ext” –global-option=”–cuda_ext .” “ Note that the extension requires CUDA>=11. Otherwise, you must turn off gradient accumulation fusion.”
allreduce_dgrad (bool required) – Do the allreduce of input gradients. The allreduce is done asynchronously with the computation of weight gradients. If sequence_parallel is True, this must be False, as no all reduce is performed.
sequence_parallel (bool required) – Indicates that sequence parallelism is used and thus in the forward pass the input is all gathered, and the backward pass the input gradients are reduce scattered.
tp_group (torch.distributed.ProcessGroup required) – The process group to use for tensor parallel operations.
grad_output_buffer (List[torch.Tensor] optional) – Buffer used to save output gradients when embedding table wgrad compute is deferred. Defaults to None.
wgrad_deferral_limit (int optional) – Limit on the number of micro-batches for which embedding weight gradient GEMM should be deferred. Disable by setting this to 0. Defaults to 0.
output_dtype (torch.dtype optional) – Optional GEMM output dtype. A dtype different from the input dtype requires Transformer Engine general_gemm.

class core.tensor_parallel.layers.ColumnParallelLinear( input_size, output_size, *, config: megatron.core.model_parallel_config.ModelParallelConfig, init_method: Callable, bias=True, gather_output=False, stride=1, keep_master_weight_for_test=False, skip_bias_add=False, skip_weight_param_allocation: bool = False, embedding_activation_buffer: Optional[List[torch.Tensor]] = None, grad_output_buffer: Optional[List[torch.Tensor]] = None, is_expert: bool = False, tp_comm_buffer_name: Optional[str] = None, disable_grad_reduce: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, output_dtype: Optional[torch.dtype] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

Bases: torch.nn.Module

Linear layer with column parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its second dimension as A = [A_1, …, A_p].

Parameters:

input_size – first dimension of matrix A.
output_size – second dimension of matrix A.
bias – If true, add bias
gather_output – If true, call all-gather on output and make Y available to all GPUs, otherwise, every GPU will have its output which is Y_i = XA_i
init_method – method to initialize weights. Note that bias is always set to zero.
stride – For the strided linear layers.
keep_master_weight_for_test – This was added for testing and should be set to False. It returns the master weights used for initialization.
skip_bias_add – If True, do not add the bias term, instead return it to be added by the caller. This enables performance optimizations where bias can be fused with other elementwise operations.
skip_weight_param_allocation – If True, weight parameter is not allocated and must be passed as a keyword argument weight during the forward pass. Note that this does not affect bias, which will be allocated if bias is True. Defaults to False.
embedding_activation_buffer – This buffer holds the input activations of the final embedding linear layer on the last pipeline stage when defer_embedding_wgrad_compute is enabled.
grad_output_buffer – This buffer holds the gradient outputs of the final embedding linear layer on the last pipeline stage when defer_embedding_wgrad_compute is enabled.
is_expert – If True, the layer is treated as an MoE expert layer.
config – ModelParallelConfig object
tp_comm_buffer_name – Communication buffer name is not used in non-Transformer-Engine modules.
disable_grad_reduce – If True, reduction of output gradients across tensor-parallel ranks will be disabled. Defaults to False. This feature is used by Lora Adapter in Nemo to delay and fuse reduction along with other gradients for performance optimization.
output_dtype – Optional dtype for the GEMM output. When it differs from the input dtype, Transformer Engine general_gemm is used.
pg_collection – Optional process group collection. Used to resolve the generalized tensor parallel remat group; falls back to the global parallel state when omitted.

Initialization

_forward_impl(input, weight, *args, **kwargs)#

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

Forward of ColumnParallelLinear

Parameters:

input_ – 3D tensor whose order of dimension is [sequence, batch, hidden]
weight (optional) – weight tensor to use, compulsory when skip_weight_param_allocation is True.
runtime_gather_output (bool) – Gather output at runtime. Default None means gather_output arg in the constructor will be used.

Returns:

output
bias

backward_dw() → None#

Compute weight gradients during the backward pass if delay_wgrad_compute is enabled.

Not supported - does nothing.

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#: Sharding along axis 0, bias sharded

set_extra_state(state: Any)#: Extra state is ignored

get_extra_state() → None#: Keep compatibility with TE state dict.

extra_repr() → str#: Extra context to add to the module’s string representation.

class core.tensor_parallel.layers.RowParallelLinear( input_size: int, output_size: int, *, config: megatron.core.model_parallel_config.ModelParallelConfig, init_method: Callable, bias: bool, input_is_parallel: bool, skip_bias_add: bool, stride: int = 1, keep_master_weight_for_test: bool = False, is_expert: bool = False, tp_comm_buffer_name: str | None = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

Bases: torch.nn.Module

Linear layer with row parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its second dimension. A = transpose([A_1 .. A_p]) X = [X_1, …, X_p]

Parameters:

input_size – first dimension of matrix A.
output_size – second dimension of matrix A.
bias – If true, add bias. Note that bias is not parallelized.
input_is_parallel – If true, we assume that the input is already split across the GPUs and we do not split again.
init_method – method to initialize weights. Note that bias is always set to zero.
stride – For the strided linear layers.
keep_master_weight_for_test – This was added for testing and should be set to False. It returns the master weights used for initialization.
skip_bias_add – If True, do not add the bias term, instead return it to be added by the caller. This enables performance optimizations where bias can be fused with other elementwise operations.
is_expert – If True, the layer is treated as an MoE expert layer
tp_comm_buffer_name – Communication buffer name. Not used in non-Transformer-Engine modules.
config – ModelParallelConfig object

Initialization

_forward_impl(input, weight, *args, **kwargs)#

forward(input_: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]#

Forward of RowParallelLinear

Parameters:

input_ – 3D tensor whose order of dimension is [sequence, batch, hidden]

Returns:

output
bias

backward_dw() → None#

Compute weight gradients during the backward pass if delay_wgrad_compute is enabled.

Not supported - does nothing.

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#: Sharding along axis 1, bias not sharded

set_extra_state(state: Any)#: Extra state is ignored

get_extra_state() → None#: Keep compatibility with TE state dict.

extra_repr() → str#: Extra context to add to the module’s string representation.

core.tensor_parallel.layers#

Module Contents#

Classes#

Functions#

Data#

API#

`core.tensor_parallel.layers`#