Implementation:NVIDIA TransformerEngine Ops AllGather

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch, Distributed
Last Updated	2026-02-07 14:00 GMT

Overview

Fusible operation that performs all-gather on tensors along the outer dimension, with reduce-scatter in the backward pass for distributed training.

Description

AllGather gathers tensors from all processes and concatenates them along the first dimension during the forward pass. In the backward pass, it performs the inverse operation (reduce-scatter) to distribute gradients. When the process group size is 1, both passes are no-ops. This operation is used in tensor-parallel configurations where sharded tensors need to be reconstructed.

Usage

Used in tensor-parallel and sequence-parallel configurations within the operation fuser to gather distributed tensor shards before computation.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/ops/basic/all_gather.py

Lines

1--84

Signature

class AllGather(BasicOperation):
    def __init__(self, process_group: Optional[torch.distributed.ProcessGroup] = None) -> None: ...
    def op_forward(self, ctx, input_, prev_op_grad_output_quantizer, next_op_input_quantizer) -> torch.Tensor: ...
    def op_backward(self, ctx, grad_output) -> Tuple[torch.Tensor, Tuple[()]]: ...

Import

from transformer_engine.pytorch.ops.basic.all_gather import AllGather

I/O Contract

Inputs

Name	Type	Required	Description
input_	torch.Tensor	Yes	Local tensor shard to gather
process_group	ProcessGroup	No	Distributed process group (default: world group)

Outputs

Name	Type	Description
output	torch.Tensor	Gathered tensor with first dim multiplied by world size

Usage Examples

from transformer_engine.pytorch.ops.basic.all_gather import AllGather

all_gather_op = AllGather(process_group=tp_group)
# In fuser pipeline: gathers local shards from all ranks

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements