`core.tensor_parallel.mappings`#

Module Contents#

Classes#

`_CopyToModelParallelRegion`	Pass the input to the model parallel region.
`_ReduceFromModelParallelRegion`	All-reduce the input from the model parallel region.
`_ScatterToModelParallelRegion`	Split the input and keep only the corresponding chuck to the rank.
`_GatherFromModelParallelRegion`	Gather the input from model parallel region and concatinate.
`_ScatterToSequenceParallelRegion`	Split the input and keep only the corresponding chuck to the rank.
`_GatherFromSequenceParallelRegion`	Gather the input from sequence parallel region and concatinate.
`_ReduceScatterToSequenceParallelRegion`	Reduce scatter the input from the model parallel region.
`_AllGatherFromTensorParallelRegion`	Gather the input from model parallel region and concatenate.
`_ReduceScatterToTensorParallelRegion`	Reduce scatter the input from the model parallel region.
`_AllToAll`

Functions#

`_reduce`	All-reduce the input tensor across model parallel group.
`_split_along_last_dim`	Split the tensor along its last dimension and keep the corresponding slice.
`_split_along_first_dim`	Split the tensor along its first dimension and keep the corresponding slice.
`_gather_along_last_dim`	Gather tensors and concatinate along the last dimension.
`_reduce_scatter_along_last_dim`	Reduce-scatter tensors on the last dimension.
`_gather_along_first_dim`	Gather tensors and concatenate along the first dimension.
`_reduce_scatter_along_first_dim`	Reduce-scatter the input tensor across model parallel group.
`copy_to_tensor_model_parallel_region`	Wrapper for autograd function: forward: copy, backward allreduce
`reduce_from_tensor_model_parallel_region`	Wrapper for autograd function: forward: all reduce, backward copy
`scatter_to_tensor_model_parallel_region`	Wrapper for autograd function: forward: RS, backward: AG
`gather_from_tensor_model_parallel_region`	Wrapper for autograd function: forward: AG, backward: split
`scatter_to_sequence_parallel_region`	Wrapper for autograd function: forward: split, backward: AG
`gather_from_sequence_parallel_region`	Wrapper for autograd function: forward: AG, backward: RS
`reduce_scatter_to_sequence_parallel_region`	Wrapper for autograd function: forward: RS, backward AG
`all_gather_last_dim_from_tensor_parallel_region`	Wrapper for autograd function: forward: AG, backward RS
`reduce_scatter_last_dim_to_tensor_parallel_region`	Wrapper for autograd function: forward: RS, backward AG: AG
`all_to_all`	Wrapper for autograd function
`all_to_all_sp2hp`	Perform AlltoAll communication on tensor parallel group, transform the input tensor from shape [num_tokens/TP, H] to [num_tokens, H/TP].
`all_to_all_hp2sp`	Perform AlltoAll communication on tensor parallel group, transform the input tensor from shape [num_tokens, H/TP] to [num_tokens/TP, H].

API#

core.tensor_parallel.mappings._reduce(input_, group)#: All-reduce the input tensor across model parallel group.

core.tensor_parallel.mappings._split_along_last_dim(input_, group)#: Split the tensor along its last dimension and keep the corresponding slice.

core.tensor_parallel.mappings._split_along_first_dim(input_, group)#: Split the tensor along its first dimension and keep the corresponding slice.

core.tensor_parallel.mappings._gather_along_last_dim(input_, group)#: Gather tensors and concatinate along the last dimension.

core.tensor_parallel.mappings._reduce_scatter_along_last_dim(input_, group)#: Reduce-scatter tensors on the last dimension.

core.tensor_parallel.mappings._gather_along_first_dim( input_, group, output_split_sizes=None, use_global_buffer=False, )#

Gather tensors and concatenate along the first dimension.

Parameters:

input_tensor (torch.Tensor) – A tensor to be gathered.
output_split_sizes (List[int], optional) – A list specifying the sizes of the output splits along the first dimension. If None, equal splitting is assumed. Default: None.

Returns:

Gathered tensor.

Return type:

torch.Tensor

core.tensor_parallel.mappings._reduce_scatter_along_first_dim( input_, group, input_split_sizes=None, use_global_buffer=False, )#

Reduce-scatter the input tensor across model parallel group.

Parameters:

input_ (torch.Tensor) – The input tensor to be reduce-scattered.
input_split_sizes (List[int], optional) – A list specifying the sizes of the input splits along the first dimension for each rank. If None, equal splitting is assumed. Default: None.

class core.tensor_parallel.mappings._CopyToModelParallelRegion#