Function allreduce

tensorrt_llm/functional.py:4048–4141

Add an operation that performs a collective all-reduce. Let's define 'world_size' as the length of the 'group' list. That functions creates a layer to compute the sum of 'world_size' tensors distributed amongst the 'world_size' participating ranks (one GPU per rank). The list

(
    tensor: Tensor,
    group: List[int],
    all_reduce_params: Optional[AllReduceParams] = AllReduceParams()
)

def allreduce(
    tensor: Tensor,
    group: List[int],
    all_reduce_params: Optional[AllReduceParams] = AllReduceParams()
) -> Tensor:
    '''
    Add an operation that performs a collective all-reduce.

    Let's define 'world_size' as the length of the 'group' list. That functions
    creates a layer to compute the sum of 'world_size' tensors distributed
    amongst the 'world_size' participating ranks (one GPU per rank).

    The list 'group' contains the identifiers of the ranks participating into
    the collective operation.

    The tensors in the different ranks must be 1D tensors (or views) and the output
    tensor will have that same shape. The output tensor will be replicated on
    the 'world_size' ranks.

    That operation is implemented using a plugin that wraps the NCCL all-reduce
    collective operation. See
    https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce
    for details.

    Parameters:
        tensor : Tensor
            The input tensor.

        group : List[int]
            The ranks participating into the all-reduce operation.

        strategy: AllReduceStrategy
            NCCL delegates all-reduce to NCCL while ONESHOT and TWOSHOT are custom latency-optimal algorithms.
            AUTO chooses amongst the three based on a message-size heuristic.

    Returns:
        The tensor produced by that layer.
    '''

    global allreduce_ub_counter
    allreduce_ub_counter += 1

    if all_reduce_params is None:
        all_reduce_params = AllReduceParams()
    all_reduce_params.update_strategy()

    # TODO(TRTLLM-996): remove this WAR when custom allreduce is supported
    # for encoder models in C++ runtime.
    workspace = None
    if all_reduce_params.strategy != AllReduceStrategy.NCCL and all_reduce_params.strategy != AllReduceStrategy.UB:
        if current_all_reduce_helper().workspace is None:
            all_reduce_params.strategy = AllReduceStrategy.NCCL_SYMMETRIC
        else:
            workspace = current_all_reduce_helper().workspace.trt_tensor
    if all_reduce_params.strategy == AllReduceStrategy.UB:
        tensor.mark_output("allreduce_ub_0_" + str(allreduce_ub_counter))
    dtype = default_net().plugin_config.nccl_plugin
    layer, allreduce_plg_creator, pfc = create_allreduce_plugin(