Bases: object
Computes the Cross Entropy Loss splitting the Vocab size across tensor parallel ranks. This implementation is used in both fused and unfused cross entropy implementations
Calculates cross entropy loss.
Calculates gradients.
Calculates logits_max.
Calculates predicted logits.
Prepare gradient calculation operands.
Performs cross entropy loss when logits are split across tensor parallel ranks
vocab_parallel_logits – logits split across tensor parallel ranks dimension is [sequence_length, batch_size, vocab_size/num_parallel_ranks]
target – correct vocab ids of dimseion [sequence_length, micro_batch_size]
label_smoothing – smoothing factor, must be in range [0.0, 1.0) default is no smoothing (=0.0)
Broadcast data from rank zero of each model parallel group to the members of the same model parallel group.
keys – list of keys in the data disctionary to be broadcasted
data – data dictionary of string keys and cpu tensor values.
datatype – torch data type of all tensors in data associated with keys.
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].
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 optimations 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.
Forward of ColumnParallelLinear
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.
output
bias
Keep compatibility with TE state dict.
Extra state is ignored
Sharding along axis 0, bias sharded
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.
Backward with frozen weight.
Forward with frozen weight.
Bases: torch.autograd.Function
See linear_with_grad_accumulation_and_async_allreduce
Backward.
Forward.
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]
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 optimations 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
Forward of RowParallelLinear
input – 3D tensor whose order of dimension is [sequence, batch, hidden]
output
bias
Keep compatibility with TE state dict.
Extra state is ignored
Sharding along axis 1, bias not sharded
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.
num_embeddings – vocabulary size.
embedding_dim – size of hidden state.
reduce_scatter_embeddings – Decides whether to perform ReduceScatter after embedding lookup
config – A megatron.core.ModelParallelConfig object
Forward.
input (torch.Tensor) – Input tensor.
Non-default implementation for embeddings due to allow_shape_mismatch param
Copy model parallel attributes from one tensor to another.
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.
Here, async and sync allreduce are the same. If sequence_parallel is True, this must be False, as no all reduce is performed.
parallelism is used and thus in the forward pass the input is all gathered, and the backward pass the input gradients are reduce scattered.
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.
Please use allreduce_dgrad instead.
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 calcluation 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.
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.
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.
async_grad_allreduce (bool optional) – Will be removed with 0.11.0. Please use allreduce_dgrad instead.
Returns true if the passed-in parameter is not a duplicate parameter on another TP rank.
Set default model parallel attributes if not set explicitly already.
Sets tp attributes to tensor
Wrapper for autograd function: forward: AG, backward RS <last dim>
Wrapper for autograd function
Perform AlltoAll communication on tensor parallel group, transform the input tensor from shape [num_tokens, H/TP] to [num_tokens/TP, H].
input (torch.Tensor) – The input tensor which has been distributed along the hidden dimension.
The output tensor with shape [num_tokens/TP, H].
torch.Tensor
Perform AlltoAll communication on tensor parallel group, transform the input tensor from shape [num_tokens/TP, H] to [num_tokens, H/TP].
input (torch.Tensor) – The input tensor which has been distributed along the sequence dimension.
The output tensor with shape [num_tokens, H/TP].
torch.Tensor
Wrapper for autograd function: forward: copy, backward allreduce
Wrapper for autograd function: forward: AG, backward: RS <first dim>
Wrapper for autograd function: forward: AG, backward: split <last dim>
Wrapper for autograd function: forward: all reduce, backward copy
Wrapper for autograd function: forward: RS, backward AG: AG <last dim>
Wrapper for autograd function: forward: RS, backward AG <fisrt dim>
Wrapper for autograd function: forward: split, backward: AG <last dim>
Wrapper for autograd function: forward: RS, backward: AG <last dim>
Bases: torch.autograd.Function
Checkpoint Function
This function is adapted from torch.utils.checkpoint with two main changes: 1) torch.cuda.set_rng_state is replaced with _set_cuda_rng_state 2) the states in the model parallel tracker are also properly tracked/set/reset.
Backward pass.
Forward pass.
Bases: object
Tracker for the cuda RNG states.
Using the add method, a cuda rng state is initialized based on the input seed and is assigned to name. Later, by forking the rng state, we can perform operations and return to our starting cuda state.
Track the rng state.
Fork the cuda rng state, perform operations, and exit with the original state.
Get rng states. Copy the dictionary so we have direct pointers to the states, not just a pointer to the dictionary.
Checks if the internal RNG state has been set wirth set_states().
Set to the initial state (no tracker).
Set the rng states. For efficiency purposes, we do not check the size of seed for compatibility.
Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.
Get cuda rng tracker.
Get the data parallel rng tracker name
Get the expert parallel rng tracker name
Create the RNG tracker. ‘use_te_rng_tracker’ determines whether to use Megatron or TransformerEngine’s implementation. In particular, TransformerEngine’s implementation is cudagraphable and supports FP8.
Initialize model parallel cuda seed.
This function should be called after the model parallel is initialized. Also, no torch.cuda.manual_seed should be called after this function. Basically, this is replacement for that function. Three set of RNG states are tracked: default state: This is for data parallelism and is the same among a set of model parallel GPUs but different across different model parallel groups. This is used for example for dropout in the non-tensor-model-parallel regions. tensor-model-parallel state: This state is different among a set of model parallel GPUs, but the same across data parallel groups. This is used for example for dropout in model parallel regions. expert-parallel-seed: This state is only used for the expert layer of MoE models. It is different among expert-tensor and expert-model parallel GPUs, and the same across expert-data parallel groups.
Bases: object
Split the vocabulary into world_size chunks and return the first and last index of the vocabulary belonging to the rank partition: Note that indices in [fist, last)
Vocab range from global vocab size.
Vocab range from per partition vocab size.
Opposite of split_tensor_into_1d_equal_chunks. Gather values from tensor model parallel ranks.
Returns a new Tensor with the gathered data.
tensor – A Tensor or view of this rank’s portion of the data.
Split a tensor along its last dimension.
tensor – input tensor.
num_partitions – number of partitions to split the tensor
contiguous_split_chunks – If True, make each chunk contiguous in memory.
A list of Tensors
Break a tensor into equal 1D chunks across tensor parallel ranks.
Returns a Tensor or View with this rank’s portion of the data.
tensor – The tensor to split
new_buffer (bool) – If True, returns a new Tensor. If False, returns a view into the existing Tensor. Default is False
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].
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 optimations 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.
Forward of ColumnParallelLinear
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.
output
bias
Keep compatibility with TE state dict.
Extra state is ignored
Sharding along axis 0, bias sharded
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]
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 optimations 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
Forward of RowParallelLinear
input – 3D tensor whose order of dimension is [sequence, batch, hidden]
output
bias
Keep compatibility with TE state dict.
Extra state is ignored
Sharding along axis 1, bias not sharded
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.
num_embeddings – vocabulary size.
embedding_dim – size of hidden state.
reduce_scatter_embeddings – Decides whether to perform ReduceScatter after embedding lookup
config – A megatron.core.ModelParallelConfig object
Forward.
input (torch.Tensor) – Input tensor.
Non-default implementation for embeddings due to allow_shape_mismatch param
Broadcast data from rank zero of each model parallel group to the members of the same model parallel group.
keys – list of keys in the data disctionary to be broadcasted
data – data dictionary of string keys and cpu tensor values.
datatype – torch data type of all tensors in data associated with keys.
Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.
Copy model parallel attributes from one tensor to another.
Wrapper for autograd function: forward: copy, backward allreduce
Wrapper for autograd function: forward: AG, backward: RS <first dim>
Wrapper for autograd function: forward: AG, backward: split <last dim>
Opposite of split_tensor_into_1d_equal_chunks. Gather values from tensor model parallel ranks.
Returns a new Tensor with the gathered data.
tensor – A Tensor or view of this rank’s portion of the data.
Get cuda rng tracker.
Get the expert parallel rng tracker name
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 calcluation 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.
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.
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.
async_grad_allreduce (bool optional) – Will be removed with 0.11.0. Please use allreduce_dgrad instead.
Initialize model parallel cuda seed.
This function should be called after the model parallel is initialized. Also, no torch.cuda.manual_seed should be called after this function. Basically, this is replacement for that function. Three set of RNG states are tracked: default state: This is for data parallelism and is the same among a set of model parallel GPUs but different across different model parallel groups. This is used for example for dropout in the non-tensor-model-parallel regions. tensor-model-parallel state: This state is different among a set of model parallel GPUs, but the same across data parallel groups. This is used for example for dropout in model parallel regions. expert-parallel-seed: This state is only used for the expert layer of MoE models. It is different among expert-tensor and expert-model parallel GPUs, and the same across expert-data parallel groups.
Returns true if the passed-in parameter is not a duplicate parameter on another TP rank.
Wrapper for autograd function: forward: all reduce, backward copy
Wrapper for autograd function: forward: RS, backward AG <fisrt dim>
Wrapper for autograd function: forward: split, backward: AG <last dim>
Wrapper for autograd function: forward: RS, backward: AG <last dim>
Set default model parallel attributes if not set explicitly already.
Sets tp attributes to tensor
Split a tensor along its last dimension.
tensor – input tensor.
num_partitions – number of partitions to split the tensor
contiguous_split_chunks – If True, make each chunk contiguous in memory.
A list of Tensors
Break a tensor into equal 1D chunks across tensor parallel ranks.
Returns a Tensor or View with this rank’s portion of the data.
tensor – The tensor to split
new_buffer (bool) – If True, returns a new Tensor. If False, returns a view into the existing Tensor. Default is False
Performs cross entropy loss when logits are split across tensor parallel ranks
vocab_parallel_logits – logits split across tensor parallel ranks dimension is [sequence_length, batch_size, vocab_size/num_parallel_ranks]
target – correct vocab ids of dimseion [sequence_length, micro_batch_size]
label_smoothing – smoothing factor, must be in range [0.0, 1.0) default is no smoothing (=0.0)
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4